Skip to content

Lazy neighborhood filter kernel compilation #1708

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
cmdupuis3 wants to merge 3 commits into
base: main
Choose a base branch
from
+104 −32
Open

Lazy neighborhood filter kernel compilation #1708

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ec83767
Lazy neighborhood filter kernel compilation
cmdupuis3 Aug 25, 2026
8564410
Lazy nb kernels with functools cache instead
cmdupuis3 Aug 25, 2026
7127726
Move kernels inside
cmdupuis3 Aug 26, 2026
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
30 changes: 30 additions & 0 deletions test/test_dependencies.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,4 +38,34 @@ def test_hvplot_optional():
_assert_not_imported_after_import_uxarray("hvplot")


def test_no_numba_kernels_built_on_import():
"""Test that `import uxarray` does not build any numba kernel.

``guvectorize`` compiles at decoration time when it is given explicit
signatures, so a kernel assigned at module scope is built during the
import. This compilation can dominate the uxarray import, and building a
``target="parallel"`` kernel starts numba's threading layer, which
leaves a thread pool running, making forks unsafe.
"""
code = (
"import numba, uxarray\n"
"try:\n"
" layer = numba.threading_layer()\n"
"except ValueError:\n"
" pass\n"
"else:\n"
" raise AssertionError(\n"
" f'`import uxarray` started numba threading layer {layer!r}. '\n"
" 'Something it imports builds a parallel kernel at module '\n"
" 'scope; build it on first use instead.'\n"
" )\n"
)
result = subprocess.run(
[sys.executable, "-c", code],
capture_output=True,
text=True,
)
assert result.returncode == 0, result.stderr


# TODO: similar tests for cartopy, holoviews, and other optional deps.
93 changes: 66 additions & 27 deletions uxarray/grid/neighbors.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import functools
import warnings
from typing import Callable

Expand Down Expand Up @@ -1277,24 +1278,62 @@ def _median(window, _):
return np.median(window)


# One compiled kernel per reduction. The methods on ``Neighborhood`` below name
# these directly, so there is no dispatch table between the public API and the
# gufuncs: a reduction is reachable only if a method exists for it, and a method
# can only reach the kernel it names. ``Neighborhood`` is the only class that
# names them -- the data-bound classes reach a kernel by naming the
# ``Neighborhood`` method for it, so there is one place per reduction where its
# kernel and parameter are chosen.
_MEAN_KERNEL = _make_kernel(lambda window, _: np.mean(window))
_SUM_KERNEL = _make_kernel(lambda window, _: np.sum(window))
_MIN_KERNEL = _make_kernel(lambda window, _: np.min(window))
_MAX_KERNEL = _make_kernel(lambda window, _: np.max(window))
_PTP_KERNEL = _make_kernel(lambda window, _: np.max(window) - np.min(window))
_MEDIAN_KERNEL = _make_kernel(_median)
_VAR_KERNEL = _make_kernel(_variance)
_STD_KERNEL = _make_kernel(lambda window, ddof: np.sqrt(_variance(window, ddof)))
# One compiled kernel per reduction. The methods on ``Neighborhood`` below call
# into these directly. Non-compiled functions are only provided hooks through
# ``Neighborhood.reduce``. If new compiled reductions are desired, they should
# follow this pattern.
#
# ``functools.cache`` defers each build to the first call. The deferred compilation
# ensures that these kernels will only be compiled individually and lazily. Further,
# the lazy compilation prevents gufuncs from spawning threadpools eagerly and
# disrupting threading and forking in other contexts.


@functools.cache
def _mean_kernel():
return _make_kernel(lambda window, _: np.mean(window))


@functools.cache
def _sum_kernel():
return _make_kernel(lambda window, _: np.sum(window))


@functools.cache
def _min_kernel():
return _make_kernel(lambda window, _: np.min(window))


@functools.cache
def _max_kernel():
return _make_kernel(lambda window, _: np.max(window))


@functools.cache
def _ptp_kernel():
return _make_kernel(lambda window, _: np.max(window) - np.min(window))


@functools.cache
def _median_kernel():
return _make_kernel(_median)


@functools.cache
def _var_kernel():
return _make_kernel(_variance)


@functools.cache
def _std_kernel():
return _make_kernel(lambda window, ddof: np.sqrt(_variance(window, ddof)))


# ``percentile`` is ``quantile`` on a 0-100 scale, so both methods rescale onto
# this one kernel rather than compiling a near-duplicate.
_QUANTILE_KERNEL = _make_kernel(lambda window, q: np.quantile(window, q))
@functools.cache
def _quantile_kernel():
return _make_kernel(lambda window, q: np.quantile(window, q))


def _as_quantile(q, scale: float):
Expand Down Expand Up @@ -1509,45 +1548,45 @@ def __repr__(self) -> str:

def mean(self, uxda):
"""Mean of each neighborhood."""
return self._apply_kernel(uxda, _MEAN_KERNEL, 0.0)
return self._apply_kernel(uxda, _mean_kernel, 0.0)

@Sevans711 Sevans711 Aug 26, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm pretty sure this isn't equivalent anymore; the equivalent syntax if sticking with functools here would be _mean_kernel(), right?

I would have a slight preference for a solution that doesn't use functools.cache, to avoid this confusion. Something like this would be more readable and less likely to cause typo, in my opinion:

_NUMBA_KERNELS = {}
def _numba_kernel(kernel_name):
    """returns numba kernel object, such as _make_kernel(lambda window, _: np.mean(window)).
    kernels are cached after initially being created.
    Supported kernels are: mean, sum, ...
    """
    if kernel_name in _NUMBA_KERNELS:
        return _NUMBA_KERNELS[kernel_name]
    # else, need to make kernel for the first time and cache it:
    if kernel_name=='mean':
        kernel = _make_kernel(lambda window, _: np.mean(window))
    elif kernel_name=='sum':
        kernel = _make_kernel(lambda window, _: np.sum(window))
    # ... similar pattern for all supported names
    else:
        raise ValueError(f'unrecognized kernel_name: {kernel_name!r}')
    _NUMBA_KERNELS[kernel_name] = kernel
    return kernel

Sidenote: upon initially reading the code I actually had the feeling that something like _NUMBA_KERNELS probably should actually be attached to the Neighborhood class for now, since these kernels are specifically used by the Neighborhood class and nowhere else. If you think these will always just be specific to the Neighborhood class, I'd recommend something like:

class Neighborhood():
    _NUMBA_KERNELS = {}

    @classmethod
    def _numba_kernel(cls, kernel_name):
        # same as above, but use cls._NUMBA_KERNELS instead.

@cmdupuis3 cmdupuis3 Aug 26, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So, the first commit on this branch took an OO approach that should be equivalent to the current approach, but I thought it was messier than the functools.cache implementation here. I think this approach has some of the same clarity issues, where we need a lot of OO and control flow to get a relatively simple result.

I'll move the kernels inside Neighborhood, and go from there.


def sum(self, uxda):
"""Sum of each neighborhood."""
return self._apply_kernel(uxda, _SUM_KERNEL, 0.0)
return self._apply_kernel(uxda, _sum_kernel, 0.0)

def min(self, uxda):
"""Smallest value in each neighborhood."""
return self._apply_kernel(uxda, _MIN_KERNEL, 0.0)
return self._apply_kernel(uxda, _min_kernel, 0.0)

def max(self, uxda):
"""Largest value in each neighborhood."""
return self._apply_kernel(uxda, _MAX_KERNEL, 0.0)
return self._apply_kernel(uxda, _max_kernel, 0.0)

def ptp(self, uxda):
"""Peak-to-peak spread (``max - min``) of each neighborhood."""
return self._apply_kernel(uxda, _PTP_KERNEL, 0.0)
return self._apply_kernel(uxda, _ptp_kernel, 0.0)

def median(self, uxda):
"""Median of each neighborhood."""
return self._apply_kernel(uxda, _MEDIAN_KERNEL, 0.0)
return self._apply_kernel(uxda, _median_kernel, 0.0)

def var(self, uxda, ddof: int = 0):
"""Variance of each neighborhood, with ``ddof`` delta degrees of
freedom."""
return self._apply_kernel(uxda, _VAR_KERNEL, float(ddof))
return self._apply_kernel(uxda, _var_kernel, float(ddof))

def std(self, uxda, ddof: int = 0):
"""Standard deviation of each neighborhood, with ``ddof`` delta degrees
of freedom."""
return self._apply_kernel(uxda, _STD_KERNEL, float(ddof))
return self._apply_kernel(uxda, _std_kernel, float(ddof))

def quantile(self, uxda, q: float):
"""Quantile ``q`` (between 0 and 1) of each neighborhood."""
return self._apply_kernel(uxda, _QUANTILE_KERNEL, _as_quantile(q, 1.0))
return self._apply_kernel(uxda, _quantile_kernel, _as_quantile(q, 1.0))

def percentile(self, uxda, q: float):
"""Percentile ``q`` (between 0 and 100) of each neighborhood."""
return self._apply_kernel(uxda, _QUANTILE_KERNEL, _as_quantile(q, 100.0))
return self._apply_kernel(uxda, _quantile_kernel, _as_quantile(q, 100.0))

def reduce(self, uxda, func: Callable):
"""Reduces each neighborhood with an arbitrary callable.
Expand Down Expand Up @@ -1592,7 +1631,7 @@ def run(block, arrays):
# path does too by writing into a float64 output.
if block.dtype not in (np.float64, np.float32):
block = block.astype(np.float64)
return kernel(block, *arrays, param)
return kernel()(block, *arrays, param)

@Sevans711 Sevans711 Aug 26, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, I see now that you added an extra call here… probably to deal with the issue I noted in my previous comment? I think this is a confusing abstraction though; basically, with this syntax it means that kernel isn't actually a compiled kernel at all, it is a "function factory which returns a compiled kernel function". This feels like it is returning to functional programming abstractions which I recall you mentioning you wanted to avoid.

If you dislike my previous suggestion and really want to keep the functools solution, I would really want to see the calls used above, e.g. self._apply_kernel(uxda, _mean_kernel(), 0.0) instead of here.

(If you really like the code as-is, the minimal change I would want to see would be to add clearer comments / docstrings to clarify for future developers that kernel isn't actually the kernel, but rather a function which returns a kernel function.)

@cmdupuis3 cmdupuis3 Aug 26, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That would be okay, either way is fine.

To be specific, I was talking about getting away from the functional-style API, rather than necessarily the underlying mechanics. I think considering how numba and @guvectorize are interacting here, a functional approach at some level is unavoidable.


return self._apply(uxda, run)

Expand Down