Antialiased lines for categorical aggregates #1081
Conversation
ianthomas23
changed the title
|
@ianthomas23 Does this summarise the issue you are seeing? from numba import njit
import numpy as np
@njit
def foo(*aggs_and_cols):
i = 0
agg = aggs_and_cols[0] # This agg is 3D
if aggs_and_cols[0].ndim == 3: # This could be any predicate
cat_index = aggs_and_cols[1][i]
agg = agg[:, :, cat_index] # This agg is 2D
xmax = agg.shape[1] - 1 # What's this agg? 2D or 3D?
args = (np.zeros((3, 4, 5)), (0, 0))
foo(*args)The problem is that Numba cannot work out the type of |
Thanks! If so, the usual way we've dealt with that in Datashader is to make sure the branch is taken before reaching Numba, so that Numba always sees the same type. |
Yes it does, thank you. Elsewhere in the code we have multiple Numba-compiled functions, each of which accepts only 2D or 3D arrays, and we pass the appropriate one of these to low-level functions like this one as an argument. The choice of which function is determined a long way up the call stack, presumably before any Numba code. So that makes sense now. |
|
Numba could in theory deal with this use case, it's just not wired up to do it yet. I'll open an issue about it because it's something that's feasible. To get this working now, two options, they are a bit contrived but should let compilation succeed.
from numba import njit
import numpy as np
@njit
def bar(aggs_and_cols0, aggs_and_cols, i):
# Numba knows how to make this into a constant and then "prune" one of the
# branches.
if aggs_and_cols0.ndim == 3:
print("3D")
tmp = aggs_and_cols0
cat_index = aggs_and_cols[1][i]
agg = tmp[:, :, cat_index]
return agg
else:
print("2D")
return aggs_and_cols0
@njit
def foo(*aggs_and_cols):
i = 0
agg = bar(aggs_and_cols[0], aggs_and_cols, i)
xmax = agg.shape[1] - 1
# 3D agg
args = (np.zeros((3, 4, 5)), (0, 0))
foo(*args)
# 2D agg
args = (np.zeros((10, 20)), (0, 0))
foo(*args)
from numba import njit
from numba.extending import overload
import numpy as np
def agg_dispatch(agg):
# write a python impl here if you want it to work with NUMBA_DISABLE_JIT=1
pass
@overload(agg_dispatch)
def ol_agg_dispatch(aggs_and_cols, i):
# In this part, you have access to the type information about aggs_and_cols
# Typically, Numba implementations would check the types to make sure they
# are appropriate for the implementations they are going to dispatch to.
agg = aggs_and_cols[0]
if agg.ndim == 2:
def impl(aggs_and_cols, i): # i is unused, just ignore it
print("2d agg")
agg = aggs_and_cols[0]
return agg
return impl
elif agg.ndim == 3:
def impl(aggs_and_cols, i):
print("3d agg")
tmp = aggs_and_cols[0]
cat_index = aggs_and_cols[1][i]
agg = tmp[:, :, cat_index]
return agg
return impl
else:
raise TypeError("Not supported")
@njit
def foo(*aggs_and_cols):
i = 0
agg = agg_dispatch(aggs_and_cols, i)
xmax = agg.shape[1] - 1
# 3D agg
args = (np.zeros((3, 4, 5)), (0, 0))
foo(*args)
# 2D agg
args = (np.zeros((10, 20)), (0, 0))
foo(*args) |
|
@stuartarchibald Thankyou, that is really useful. |
No problem. I've opened numba/numba#8021 to track, in theory I think this can be handled. |
|
I have implemented @stuartarchibald's second fix. It works really well, and the changes are localised so they cannot affect anything outside of the antialiased line code. Test added. I'm removing the WIP label as this is ready to go if the CI passes. |
ianthomas23
changed the title
Codecov Report
@@ Coverage Diff @@
## master #1081 +/- ##
==========================================
+ Coverage 83.31% 83.34% +0.03%
==========================================
Files 34 34
Lines 7473 7488 +15
==========================================
+ Hits 6226 6241 +15
Misses 1247 1247
Continue to review full report at Codecov.
|
|
I have excluded the new |
Great, glad the fix worked! |
This is a WIP to fix issue #1079. It works with
numbadisabled but not withnumbaenabled so needs further work.Test code:
When jupyter lab is run using
NUMBA_DISABLE_JIT=1 jupyter labit gives the desired output (left image uses non-antialiased code, the other 2 images use the new antialiased code):If run without disabling numba the initial failure is (full traceback is at the bottom of this post):
The
aggargument to this function is a 2Dnumpyarray for non-categorical aggregations, which works fine in themasterbranch, and a 3Dnumpyfor categorical aggregations (the length of the third axis is the number of categories).I believe that the prior code deals with this by passing down an appropriate numba-compiled
appendfunction created here:datashader/datashader/compiler.py
Line 100 in a7bc8ec
which access the
aggand up until now this has not been necessary for the new antialiasing code.It would be great if someone has some insight into how we can get this working in the short term, even if the solution is inelegant. In the medium term I have volunteered to become proficient with
numbaso that I can refactor datashader's numba code in the belief that it can be simplified as numba is now much cleverer than when this was written n years ago.Full traceback: