CODEWITHSUNDEEP
awkward-array
DragonTux
DragonTux

Reputation: 742

Implementing numpy.linalg.norm for awkward-arrays

At the moment, no implementation exists for numpy.linalg.norm with akward-arrays. In awkward1 I can use behavior to overwrite (or add?) implementations as shown in the tutorial. In my use-case I have successfully implemented versions of np.absolute and np.subtract for TVector2 and TVector3 using the tutorial example. However, numpy.linalg.norm fails.

How can I define a behavior for awkward-arrays? E.g. in the example below I tried this with a simple int64 array:

import awkward1 as ak
import numpy as np

def norm(data, *args, **kwargs):
    return np.absolute(data)

ak.behavior[np.linalg.norm, "int64"] = norm

sample = ak.Array(np.arange(10))

print(ak.type(sample))  # 10 * int64
print(np.absolute(sample))  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

print(np.linalg.norm(np.arange(10)))  # 16.881943016134134
print(np.linalg.norm(sample)) # raises TypeError

Upvotes: 1

Views: 350

Answers (1)

Jim Pivarski
Jim Pivarski

Reputation: 5974

If you haven't seen https://awkward-array.readthedocs.io/en/latest/ak.behavior.html , it's the place to start.

If np.linalg.norm is a ufunc, the following should work:

ak.behavior[np.linalg.norm, "vector2"] = lambda a: np.sqrt(a.x**2 +a.y**2)

where the data are in records with fields x and y that are named "vector2". You can make such a thing with

ak.zip({"x": x_array, "y": y_array}, with_name="vector2")

(Caveat: I haven't tested this—I wrote it on my phone. I'll test it later and edit this answer, or you might be able to edit it.)

If it's not a ufunc, then there isn't a publicly accessible (i.e. no underscores) way to define it, and perhaps there should be.

Also, be aware of the https://github.com/scikit-hep/vector project, which is defining operations like this.

Edit: Looking at the np.linalg.norm documentation, you're right, it's not a ufunc. (isinstance(np.linalg.norm, np.ufunc) is False.)

Also, this function makes some strong assumptions about its input. Two-dimensional arrays are interpreted as matrices, rather than lists of vectors. In Awkward Array, such distinctions would be labeled by metadata. I don't think we want to implement a function like this, since it could wreak havoc on jagged arrays that are supposed to be lists of different length vectors.

I think what you want is to add vector behaviors, like this:

>>> class Vector2(ak.Record):
...     @property
...     def norm(self):
...         return np.sqrt(self.x**2 + self.y**2)
... 
>>> class ArrayOfVector2(ak.Array):
...     @property
...     def norm(self):
...         return np.sqrt(self.x**2 + self.y**2)
... 
>>> ak.behavior["vector2"] = Vector2
>>> ak.behavior["*", "vector2"] = ArrayOfVector2

This just says that if you have a record that is named "vector2", then it should be presented as Python class Vector2, and if you have an array (any depth) of "vector2" records, they should collectively be presented as an ArrayOfVector2 class.

>>> x_array = ak.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]])
>>> y_array = ak.Array([[1.0, 2.0, 3.0], [], [4.0, 5.0]])
>>> vectors = ak.zip({"x": x_array, "y": y_array}, with_name="vector2")
>>> vectors[0, 0]
<Vector2 {x: 1.1, y: 1} type='vector2["x": float64, "y": float64]'>
>>> vectors
<ArrayOfVector2 [[{x: 1.1, y: 1}, ... y: 5}]] type='3 * var * vector2["x...'>

That lets you use the methods and properties you've defined:

>>> vectors[0, 0].norm
1.4866068747318506
>>> vectors.norm
<Array [[1.49, 2.97, 4.46], ... [5.95, 7.43]] type='3 * var * float64'>

But also, this is exactly what the https://github.com/scikit-hep/vector project is doing. Keep an eye on it!

Upvotes: 1

Related Questions