Feature request for typetracers handler for np.interp
#493
Comments
|
@green-cabbage it's probably something along the lines of: def interp(vals, knots_x, knots_y):
result = ak.Array(np.interp(
ak.typetracer.length_zero_if_typetracer(vals),
ak.typetracer.length_zero_if_typetracer(knots_x),
ak.typetracer.length_zero_if_typetracer(knots_y),
))
if ak.backend(vals, knots_x, knots_y) == "typetracer":
return ak.Array(result.layer.to_typetracer(forget_length=True))
return result
interpolated = dak.map_partitions(
interp,
vals,
xs,
ys,
)If you have jagged arrays you'll have to do the flattening and unflattening as well. |
|
Here's a quick code that I wrote up to test this: This returns |
|
You left out all the length_zero_if_typetracers that are necessary. |
|
You cannot mix delayed mode evaluation with eager mode evaluation. |
|
i.e. if you have an interpolation lookup table to implement you need to ship that with the function. You can make a class that has a it'll be something like: class MyScaleVariation:
def __init__(self, x_knots, y_knots):
self.x_knots = x_knots
self.y_knots = y_knots
def __call__(self, vals):
# handle the typetracers
out = np.iterp(vals, self.x_knots, self.y_knots)
return outThen you make that class, and then you call map_partitions on it, or if you want to get fancy you can check for a dask awkward array being passed into the |
|
I have taken your suggestion and made a new code: This goes further, but then returns this error: |
|
OK I need to be far more precise for you to get something that works. You are shooting in the dark here a bit, and I think it is more the documentation's fault than anything else. The lifecycle of evaluating functions you send to map_partitions is:
class MyScaleVariation:
def __init__(self, x_knots, y_knots):
self.x_knots = x_knots
self.y_knots = y_knots
def __call__(self, vals):
result = ak.Array(np.interp(
ak.typetracer.length_zero_if_typetracer(vals), # this will either be a concrete array with data or a type tracer
self.x_knots,
self.y_knots,
))
print(f"result: {result}")
if ak.backend(vals) == "typetracer":
return ak.Array(result.layout.to_typetracer(forget_length=True))
return result
scalevar = MyScaleVariation(xp, fp) # xp and fp should be *concrete* arrays that you use later.
interpolated = dak.map_partitions(
scalevar,
jets.pt.flatten() # just make sure you can unflatten it too, if you need that.
) |
|
Your last suggestion essentially solved my issue. My code below functions as intended: This prints an actual array, which is great! I will now try to integrate this into my framework, and will let you know if things don't go well. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Hello, I originally made a QA question on implementing dask_awkward equivalent of
np.interpin this QA discussion.@lgray suggested wrapping
np.interpwithdak.map_partitions, but doing so gaveTypeError("cannot compare unknown lengths against known values"). The full error log can be seen in here. @lgray later said that this indicates that some wrapper code is neceesary to handle the typetracers. Under his suggestion, I would like to make a feature request for this functionality, so that I can in turn wrapnp.interpwithdak.map_partitions. Please let me know if you need more info!The text was updated successfully, but these errors were encountered: