Fix typing errors when upgrading to NumPy 1.20.0

sgkit/distance/api.py

@@ -100,4 +100,4 @@ def pairwise_distance(
         concatenate=True,
     )
     x_distance = da.triu(x_distance, 1) + da.triu(x_distance).T
-    return x_distance.compute()
+    return x_distance.compute()  # type: ignore[no-any-return]

sgkit/io/bgen/bgen_reader.py

@@ -145,7 +145,7 @@ def __getitem__(self, idx: Any) -> np.ndarray:
                     res = np.zeros((len(all_vaddr), len(probs), 3), dtype=self.dtype)
                 res[i] = probs
             res = res[..., idx[2]]  # type: ignore[index]
-            return np.squeeze(res, axis=squeeze_dims)
+            return np.squeeze(res, axis=squeeze_dims)  # type: ignore[no-any-return]
 
 
 def _split_alleles(allele_ids: bytes) -> List[bytes]:
@@ -362,7 +362,7 @@ def encode_variables(
             # 16 bits will cause overflow/underflow
             # See https://en.wikipedia.org/wiki/Floating-point_arithmetic#Internal_representation
             # *bits precision column for single precision floats
-            if dtype not in [np.uint8, np.uint16]:
+            if dtype not in [np.uint8, np.uint16]:  # type: ignore[comparison-overlap]
                 raise ValueError(
                     "Probability integer dtype invalid, must "
                     f"be uint8 or uint16 not {probability_dtype}"

sgkit/io/plink/plink_reader.py

@@ -84,7 +84,7 @@ def __getitem__(self, idx: Tuple[Any, ...]) -> np.ndarray:
         call1 = np.where(arr < 0, -1, np.where(arr == 2, 1, 0))
         arr = np.stack([call0, call1], axis=-1)
         # Apply final slice to 3D result
-        return arr[:, :, idx[-1]]
+        return arr[:, :, idx[-1]]  # type: ignore[no-any-return]
 
     def close(self) -> None:
         # This is not actually crucial since a Bed instance with no
@@ -265,7 +265,7 @@ def read_plink(
     # Otherwise create index for contig names based
     # on order of appearance in underlying .bim file
     else:
-        variant_contig, variant_contig_names = encode_array(arr_bim["contig"].compute())
+        variant_contig, variant_contig_names = encode_array(arr_bim["contig"].compute())  # type: ignore
         variant_contig = variant_contig.astype("int16")
         variant_contig_names = list(variant_contig_names)
 

sgkit/io/utils.py

@@ -183,7 +183,7 @@ def load_chunk(
 
 def _zarr_index(offsets: ArrayLike, pos: int) -> int:
     """Return the index of the zarr file that pos falls in"""
-    index: int = np.searchsorted(offsets, pos, side="right") - 1
+    index: int = np.searchsorted(offsets, pos, side="right") - 1  # type: ignore[assignment]
     return index
 
 

sgkit/io/vcfzarr_reader.py

@@ -196,7 +196,7 @@ def _vcfzarr_to_dataset(
                 # Compute fixed-length string dtype for array
                 if kind == "O" or var in ("variant_id", "variant_allele"):
                     kind = "S"
-                max_len = max_str_len(arr).values
+                max_len = max_str_len(arr).values  # type: ignore[union-attr]
                 dt = f"{kind}{max_len}"
                 ds[var] = arr.astype(dt)  # type: ignore[no-untyped-call]
 

sgkit/stats/hwe.py

@@ -48,7 +48,7 @@ def hardy_weinberg_p_value(obs_hets: int, obs_hom1: int, obs_hom2: int) -> float
     het_probs = np.zeros(obs_mac + 1, dtype=np.float64)
 
     if obs_n == 0:
-        return np.nan  # type: ignore[no-any-return]
+        return np.nan
 
     # Identify distribution midpoint
     mid = int(obs_mac * (2 * obs_n - obs_mac) / (2 * obs_n))
@@ -91,7 +91,7 @@ def hardy_weinberg_p_value(obs_hets: int, obs_hom1: int, obs_hom2: int) -> float
         curr_hets += 2
 
     if prob_sum <= 0:  # pragma: no cover
-        return np.nan  # type: ignore[no-any-return]
+        return np.nan
     het_probs = het_probs / prob_sum
     p = het_probs[het_probs <= het_probs[obs_hets]].sum()
     p = max(min(1.0, p), 0.0)

sgkit/stats/pca.py

@@ -130,7 +130,7 @@ def pca_stats(ds: Dataset, est: BaseEstimator, *, merge: bool = True) -> Dataset
     if "sample_pca_component" in new_ds and "sample_pca_explained_variance" in new_ds:
         new_ds[variables.sample_pca_loading] = new_ds[
             variables.sample_pca_component
-        ] * np.sqrt(new_ds[variables.sample_pca_explained_variance])
+        ] * np.sqrt(new_ds[variables.sample_pca_explained_variance].data)
     return conditional_merge_datasets(ds, variables.validate(new_ds), merge)
 
 

sgkit/stats/regenie.py

@@ -598,7 +598,7 @@ def regenie_transform(
     *,
     variant_block_size: Optional[Union[int, Tuple[int, ...]]] = None,
     sample_block_size: Optional[Union[int, Tuple[int, ...]]] = None,
-    alphas: Optional[Sequence[float]] = None,
+    alphas: Optional[ArrayLike] = None,
     add_intercept: bool = True,
     orthogonalize: bool = False,
     normalize: bool = False,

sgkit/tests/test_pca.py

@@ -29,7 +29,7 @@ def simulate_cohort_genotypes(
         ]
     )
     # Sample allele counts in [0, 1, 2]
-    return rs.binomial(2, af.T).astype("int8")
+    return rs.binomial(2, af.T).astype("int8")  # type: ignore[no-any-return]
 
 
 def simulate_dataset(

sgkit/tests/test_preprocessing.py

@@ -85,7 +85,7 @@ def test_patterson_scaler__missing_data(dtype, use_nan):
         ac = np.where(ac < 0, np.nan, ac)
     scaler = sgkit.stats.preprocessing.PattersonScaler().fit(ac)
     # Read means from columns of array; scale = sqrt(mean/2 * (1 - mean/2))
-    np.testing.assert_equal(scaler.mean_, np.array([1] * 3 + [np.nan]))
+    np.testing.assert_equal(scaler.mean_, np.array([1.0] * 3 + [np.nan]))
     np.testing.assert_equal(scaler.scale_, np.array([0.5] * 3 + [np.nan]))
     np.testing.assert_equal(
         scaler.transform(ac),
@@ -121,7 +121,7 @@ def test_filter_partial_calls():
 
     np.testing.assert_array_equal(
         mask_filtered,
-        np.array([[[0, 0], [0, 0], [0, 0]], [[1, 1], [1, 1], [1, 1]]], dtype=np.bool),
+        np.array([[[0, 0], [0, 0], [0, 0]], [[1, 1], [1, 1], [1, 1]]], dtype=bool),
     )
 
 
@@ -159,7 +159,7 @@ def test_filter_partial_calls__mixed_ploidy():
                 [[0, 0, 0, 0], [0, 0, 1, 1], [0, 0, 0, 0]],
                 [[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]],
             ],
-            dtype=np.bool,
+            dtype=bool,
         ),
     )
 

sgkit/window.py

@@ -106,7 +106,7 @@ def has_windows(ds: Dataset) -> bool:
 
 
 def moving_statistic(
-    values: ArrayLike,
+    values: da.Array,
     statistic: Callable[..., ArrayLike],
     size: int,
     step: int,