From 85be19928111d32bbbf17af72ee9af8324f11b28 Mon Sep 17 00:00:00 2001
From: Walter Simson <waltersimson@stanford.edu>
Date: Tue, 19 Nov 2024 15:29:22 -0800
Subject: [PATCH] Fix numpy warnings (#510)

* update linalg import

* explicitly cast to intp

* remove @typechecker

* cast to int16 to replicate legacy behavior

* update types for larger values

* try uint16

* convert to uint16

* revert change

* test two step

* convert uint to unsigned int to allow negative indexing instead of wraparound with unsinged ints

* fix type
---
 kwave/ktransducer.py       |  2 +-
 kwave/utils/kwave_array.py |  2 +-
 kwave/utils/matlab.py      | 11 ++++++++---
 3 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/kwave/ktransducer.py b/kwave/ktransducer.py
index 176e3f05..0f0f3315 100644
--- a/kwave/ktransducer.py
+++ b/kwave/ktransducer.py
@@ -644,7 +644,7 @@ def delay_mask(self, mode=None):
             ].min()  # -1s compatibility
         else:
             mask[unflatten_matlab_mask(mask, active_elements_index - 1)] += self.stored_beamforming_delays_offset  # -1s compatibility
-        return mask.astype(np.uint8)
+        return mask.astype(np.uint16)
 
     @property
     def elevation_beamforming_delays(self):
diff --git a/kwave/utils/kwave_array.py b/kwave/utils/kwave_array.py
index 9996e055..3d666dbb 100644
--- a/kwave/utils/kwave_array.py
+++ b/kwave/utils/kwave_array.py
@@ -6,7 +6,7 @@
 
 import numpy as np
 from numpy import arcsin, pi, cos, size, array
-from numpy.linalg import linalg
+import numpy.linalg as linalg
 
 from kwave.data import Vector
 from kwave.kgrid import kWaveGrid
diff --git a/kwave/utils/matlab.py b/kwave/utils/matlab.py
index a589380c..9c0e43c5 100644
--- a/kwave/utils/matlab.py
+++ b/kwave/utils/matlab.py
@@ -1,5 +1,5 @@
 from typing import Tuple, Union, Optional, List
-
+from beartype import beartype as typechecker
 import numpy as np
 
 
@@ -75,6 +75,7 @@ def matlab_find(arr: Union[List[int], np.ndarray], val: int = 0, mode: str = "ne
     return np.expand_dims(arr, -1)  # compatibility, n => [n, 1]
 
 
+@typechecker
 def matlab_mask(arr: np.ndarray, mask: np.ndarray, diff: Optional[int] = None) -> np.ndarray:
     """
     Applies a mask to an array and returns the masked elements.
@@ -89,10 +90,14 @@ def matlab_mask(arr: np.ndarray, mask: np.ndarray, diff: Optional[int] = None) -
 
     """
 
+    if mask.dtype == "uint8":
+        mask = mask.astype("int8")
+
     if diff is None:
-        return np.expand_dims(arr.ravel(order="F")[mask.ravel(order="F")], axis=-1)  # compatibility, n => [n, 1]
+        flat_mask = mask.ravel(order="F")
     else:
-        return np.expand_dims(arr.ravel(order="F")[mask.ravel(order="F") + diff], axis=-1)  # compatibility, n => [n, 1]
+        flat_mask = mask.ravel(order="F") + diff
+    return np.expand_dims(arr.ravel(order="F")[flat_mask], axis=-1)  # compatibility, n => [n, 1]
 
 
 def unflatten_matlab_mask(arr: np.ndarray, mask: np.ndarray, diff: Optional[int] = None) -> Tuple[Union[int, np.ndarray], ...]: