Modernize NumPy random functions, fix mypy errors for issue#756

Author: mohamed-laarej

malariagen_data/anoph/sample_metadata.py

@@ -1,6 +1,17 @@
 import io
 from itertools import cycle
-from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Union
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Mapping,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+    cast,
+)
 
 import ipyleaflet  # type: ignore
 import numpy as np
@@ -11,6 +22,7 @@
 from ..util import check_types
 from . import base_params, map_params, plotly_params
 from .base import AnophelesBase
+from numpy.typing import NDArray
 
 
 class AnophelesSampleMetadata(AnophelesBase):
@@ -891,8 +903,11 @@ def _prep_sample_selection_cache_params(
             # integer indices instead.
             df_samples = self.sample_metadata(sample_sets=sample_sets)
             sample_query_options = sample_query_options or {}
-            loc_samples = df_samples.eval(sample_query, **sample_query_options).values
-            sample_indices = np.nonzero(loc_samples)[0].tolist()
+            loc_samples = cast(
+                NDArray[Any],
+                df_samples.eval(sample_query, **sample_query_options).values,
+            )
+            sample_indices = cast(List[int], np.nonzero(loc_samples)[0].tolist())
 
         return sample_sets, sample_indices
 

malariagen_data/anoph/snp_frq.py

@@ -7,7 +7,6 @@
 from numpydoc_decorator import doc  # type: ignore
 import xarray as xr
 import numba  # type: ignore
-
 from .. import veff
 from ..util import (
     check_types,
@@ -576,8 +575,8 @@ def snp_allele_frequencies_advanced(
                 raise ValueError("No SNPs remaining after dropping invariant SNPs.")
 
             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)
+            count = np.compress(loc_variant, count, axis=0).reshape(-1, count.shape[1])
+            nobs = np.compress(loc_variant, nobs, axis=0).reshape(-1, nobs.shape[1])
             frequency = np.compress(loc_variant, frequency, axis=0)
 
         # Set up variant effect annotator.