Improving grid.vicinity_offset performance optionally with numba
This commit is contained in:
parent
4400c83da3
commit
e03ad0db11
|
@ -19,6 +19,14 @@ from . import Rotation
|
||||||
from . import Table
|
from . import Table
|
||||||
from . import Colormap
|
from . import Colormap
|
||||||
from ._typehints import FloatSequence, IntSequence, IntCollection, NumpyRngSeed
|
from ._typehints import FloatSequence, IntSequence, IntCollection, NumpyRngSeed
|
||||||
|
try:
|
||||||
|
import numba as nb # type: ignore
|
||||||
|
except ImportError:
|
||||||
|
nb = False
|
||||||
|
|
||||||
|
def numba_njit_wrapper(**kwargs):
|
||||||
|
return (lambda function: nb.njit(function) if nb else function)
|
||||||
|
|
||||||
|
|
||||||
class Grid:
|
class Grid:
|
||||||
"""
|
"""
|
||||||
|
@ -1129,7 +1137,7 @@ class Grid:
|
||||||
"""
|
"""
|
||||||
def most_frequent(stencil: np.ndarray,
|
def most_frequent(stencil: np.ndarray,
|
||||||
selection: Union[None,set],
|
selection: Union[None,set],
|
||||||
rng):
|
rng: np.random.Generator):
|
||||||
me = stencil[stencil.size//2]
|
me = stencil[stencil.size//2]
|
||||||
if selection is None or me in selection:
|
if selection is None or me in selection:
|
||||||
unique, counts = np.unique(stencil,return_counts=True)
|
unique, counts = np.unique(stencil,return_counts=True)
|
||||||
|
@ -1289,19 +1297,27 @@ class Grid:
|
||||||
Updated grid-based geometry.
|
Updated grid-based geometry.
|
||||||
|
|
||||||
"""
|
"""
|
||||||
def tainted_neighborhood(stencil: np.ndarray, selection: Union[None,set]):
|
@numba_njit_wrapper()
|
||||||
|
def tainted_neighborhood(stencil: np.ndarray,
|
||||||
|
selection: Optional[np.ndarray] = None):
|
||||||
me = stencil[stencil.size//2]
|
me = stencil[stencil.size//2]
|
||||||
return np.any(stencil != me if selection is None else
|
if selection is None:
|
||||||
np.in1d(stencil,np.array(list(selection - {me}))))
|
return np.any(stencil != me)
|
||||||
|
elif not len(selection)==0:
|
||||||
|
for stencil_item in stencil:
|
||||||
|
for selection_item in selection:
|
||||||
|
if stencil_item==selection_item and selection_item!=me:
|
||||||
|
return True
|
||||||
|
return False
|
||||||
d = np.floor(distance).astype(np.int64)
|
d = np.floor(distance).astype(np.int64)
|
||||||
ext = np.linspace(-d,d,1+2*d,dtype=float),
|
ext = np.linspace(-d,d,1+2*d,dtype=float),
|
||||||
xx,yy,zz = np.meshgrid(ext,ext,ext)
|
xx,yy,zz = np.meshgrid(ext,ext,ext)
|
||||||
footprint = xx**2+yy**2+zz**2 <= distance**2+distance*1e-8
|
footprint = xx**2+yy**2+zz**2 <= distance**2+distance*1e-8
|
||||||
offset_ = np.nanmax(self.material)+1 if offset is None else offset
|
offset_ = np.nanmax(self.material)+1 if offset is None else offset
|
||||||
selection_ = None if selection is None else \
|
selection_ = None if selection is None else \
|
||||||
set(self.material.flatten()) - set(util.aslist(selection)) if invert_selection else \
|
np.array(list(set(self.material.flatten()) - set(util.aslist(selection)))) if invert_selection else \
|
||||||
set(self.material.flatten()) & set(util.aslist(selection))
|
np.array(list(set(self.material.flatten()) & set(util.aslist(selection))))
|
||||||
|
|
||||||
mask = ndimage.generic_filter(self.material,
|
mask = ndimage.generic_filter(self.material,
|
||||||
tainted_neighborhood,
|
tainted_neighborhood,
|
||||||
footprint=footprint,
|
footprint=footprint,
|
||||||
|
|
Loading…
Reference in New Issue