Commit 537a3a5c authored by Sævik Pål Næverlid's avatar Sævik Pål Næverlid
Browse files

Get correct results for masked arrays

parent 67016f23
......@@ -879,4 +879,23 @@ def compute_w(pn, pm, u, v, z_w, z_r):
def nearest_unmasked(mask, i, j):
return i, j
# All neighbours
i_center = np.int32(np.round(i))
j_center = np.int32(np.round(j))
i_neigh_raw = i_center + np.array([0, 1, 1, 0, -1, -1, -1, 0, 1])[:, np.newaxis]
j_neigh_raw = j_center + np.array([0, 0, 1, 1, 1, 0, -1, -1, -1])[:, np.newaxis]
# Handle neighbours outside the domain
i_neigh = np.clip(i_neigh_raw, 0, mask.shape[1] - 1)
j_neigh = np.clip(j_neigh_raw, 0, mask.shape[0] - 1)
# Compute distance to origin
dist = np.abs(i_neigh - i) + np.abs(j_neigh - j)
dist_mask = np.ma.masked_array(dist, mask[j_neigh, i_neigh])
# Find coordinates of closest unmasked cell
idx = dist_mask.argmin(axis=0)
i_close = i_neigh[idx, np.arange(len(idx))]
j_close = j_neigh[idx, np.arange(len(idx))]
return i_close, j_close
......@@ -10,3 +10,11 @@ class Test_nearest_unmasked:
ii, jj = gridforce.nearest_unmasked(mask, i, j)
assert ii.tolist() == i.tolist()
assert jj.tolist() == j.tolist()
def test_correct_when_south_edge(self):
mask_south = np.array([[0, 0, 0, 0], [1, 1, 1, 1], [1, 1, 1, 1]])
i = np.array([0, 1, 2.51])
j = np.array([0, 1, 1.49])
ii, jj = gridforce.nearest_unmasked(mask_south, i, j)
assert ii.tolist() == [0, 1, 3]
assert jj.tolist() == [0, 0, 0]
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment