Skip to content

Use scipy distance_transform_edt to replace slow loop #947

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
trhille wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Use scipy distance_transform_edt to replace slow loop #947

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
114 changes: 52 additions & 62 deletions compass/landice/mesh.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
from mpas_tools.mesh.creation.sort_mesh import sort_mesh
from netCDF4 import Dataset
from scipy.interpolate import NearestNDInterpolator, interpn
from scipy.ndimage import distance_transform_edt


def mpas_flood_fill(seed_mask, grow_mask, cellsOnCell, nEdgesOnCell,
Expand Down Expand Up @@ -435,8 +436,8 @@ def set_cell_width(self, section_name, thk, bed=None, vx=None, vy=None,
return cell_width


def get_dist_to_edge_and_gl(self, thk, topg, x, y,
section_name, window_size=None):
def get_dist_to_edge_and_gl(self, thk, topg, x, y, section_name,
window_size=None):
"""
Calculate distance from each point to ice edge and grounding line,
to be used in mesh density functions in
Expand Down Expand Up @@ -486,8 +487,10 @@ def get_dist_to_edge_and_gl(self, thk, topg, x, y,
dist_to_grounding_line : numpy.ndarray
Distance from each cell to the grounding line
"""

logger = self.logger
section = self.config[section_name]

tic = time.time()

high_dist = float(section.get('high_dist'))
Expand All @@ -496,78 +499,65 @@ def get_dist_to_edge_and_gl(self, thk, topg, x, y,
if window_size is None:
window_size = max(high_dist, high_dist_bed)
elif window_size < min(high_dist, high_dist_bed):
logger.info('WARNING: window_size was set to a value smaller'
' than high_dist and/or high_dist_bed. Resetting'
f' window_size to {max(high_dist, high_dist_bed)},'
' which is max(high_dist, high_dist_bed)')
logger.info(
'WARNING: window_size was set smaller than high_dist and/or '
'high_dist_bed. Resetting window_size to '
f'{max(high_dist, high_dist_bed)}'
)
window_size = max(high_dist, high_dist_bed)

dx = x[1] - x[0] # assumed constant and equal in x and y
nx = len(x)
ny = len(y)
sz = thk.shape
dx = float(x[1] - x[0])
dy = float(y[1] - y[0])

# Create masks to define ice edge and grounding line
neighbors = np.array([[1, 0], [-1, 0], [0, 1], [0, -1],
[1, 1], [-1, 1], [1, -1], [-1, -1]])
# Same masks as the current implementation
neighbors = np.array([
[1, 0], [-1, 0], [0, 1], [0, -1],
[1, 1], [-1, 1], [1, -1], [-1, -1]
])

ice_mask = thk > 0.0
grounded_mask = thk > (-1028.0 / 910.0 * topg)
margin_mask = np.zeros(sz, dtype='i')
grounding_line_mask = np.zeros(sz, dtype='i')

margin_mask = np.zeros(thk.shape, dtype=bool)
grounding_line_mask = np.zeros(thk.shape, dtype=bool)

for n in neighbors:
not_ice_mask = np.logical_not(np.roll(ice_mask, n, axis=[0, 1]))
margin_mask = np.logical_or(margin_mask, not_ice_mask)

not_grounded_mask = np.logical_not(np.roll(grounded_mask,
n, axis=[0, 1]))
grounding_line_mask = np.logical_or(grounding_line_mask,
not_grounded_mask)

# where ice exists and neighbors non-ice locations
margin_mask = np.logical_and(margin_mask, ice_mask)
# optional - plot mask
# plt.pcolor(margin_mask); plt.show()

# Calculate dist to margin and grounding line
[XPOS, YPOS] = np.meshgrid(x, y)
dist_to_edge = np.zeros(sz)
dist_to_grounding_line = np.zeros(sz)

d = int(np.ceil(window_size / dx))
rng = np.arange(-1 * d, d, dtype='i')
max_dist = float(d) * dx
margin_mask |= not_ice_mask

# just look over areas with ice
# ind = np.where(np.ravel(thk, order='F') > 0)[0]
ind = np.where(np.ravel(thk, order='F') >= 0)[0] # do it everywhere
for iii in range(len(ind)):
[i, j] = np.unravel_index(ind[iii], sz, order='F')

irng = i + rng
jrng = j + rng

# only keep indices in the grid
irng = irng[np.nonzero(np.logical_and(irng >= 0, irng < ny))]
jrng = jrng[np.nonzero(np.logical_and(jrng >= 0, jrng < nx))]

dist_to_here = ((XPOS[np.ix_(irng, jrng)] - x[j]) ** 2 +
(YPOS[np.ix_(irng, jrng)] - y[i]) ** 2) ** 0.5

dist_to_here_edge = dist_to_here.copy()
dist_to_here_grounding_line = dist_to_here.copy()

dist_to_here_edge[margin_mask[np.ix_(irng, jrng)] == 0] = max_dist
dist_to_here_grounding_line[grounding_line_mask
[np.ix_(irng, jrng)] == 0] = max_dist

dist_to_edge[i, j] = dist_to_here_edge.min()
dist_to_grounding_line[i, j] = dist_to_here_grounding_line.min()
not_grounded_mask = np.logical_not(
np.roll(grounded_mask, n, axis=[0, 1])
)
grounding_line_mask |= not_grounded_mask

# Preserve current semantics for margin
margin_mask &= ice_mask

# NOTE:
# The current code does *not* apply "& ice_mask" to grounding_line_mask.
# If that was intentional, keep it. If it was accidental, add:
# grounding_line_mask &= ice_mask

# EDT computes distance to nearest zero.
# So invert the boundary masks: non-boundary=1, boundary=0.
dist_to_edge = distance_transform_edt(
~margin_mask, sampling=(dy, dx)
)
dist_to_grounding_line = distance_transform_edt(
~grounding_line_mask, sampling=(dy, dx)
)

# Preserve the current "window_size" behavior by clipping large distances.
# The old code returned max_dist for anything outside the search window.
max_dist = float(np.ceil(window_size / max(dx, dy))) * max(dx, dy)
dist_to_edge = np.minimum(dist_to_edge, max_dist)
dist_to_grounding_line = np.minimum(dist_to_grounding_line, max_dist)

toc = time.time()
logger.info('compass.landice.mesh.get_dist_to_edge_and_gl() took {:0.2f} '
'seconds'.format(toc - tic))
logger.info(
'compass.landice.mesh.get_dist_to_edge_and_gl() took '
f'{toc - tic:0.2f} seconds'
)

return dist_to_edge, dist_to_grounding_line

Expand Down
Loading