如果您还没有看到https://awkward-array.readthedocs.io/en/latest/ak.behavior.html,那就从这里开始吧。
如果np.linalg.norm
是一个ufunc,以下应该工作:
ak.behavior[np.linalg.norm, "vector2"] = lambda a: np.sqrt(a.x**2 +a.y**2)
其中数据位于具有字段x
且y
名为 的记录中"vector2"
。你可以用
ak.zip({"x": x_array, "y": y_array}, with_name="vector2")
(警告:我还没有测试过——我是在手机上写的。我稍后会测试并编辑这个答案,或者你可以编辑它。)
如果它不是 ufunc,那么就没有可公开访问(即没有下划线)的方式来定义它,也许应该有。
另外,请注意https://github.com/scikit-hep/vector项目,它定义了这样的操作。
编辑:查看np.linalg.norm文档,你是对的,它不是 ufunc。(isinstance(np.linalg.norm, np.ufunc)
是False
。)
此外,此函数对其输入做出了一些强有力的假设。二维数组被解释为矩阵,而不是向量列表。在 Awkward Array 中,此类区别将由元数据标记。我不认为我们想要实现这样的函数,因为它可能会对本来应该是不同长度向量列表的锯齿状数组造成严重破坏。
我认为您想要的是添加矢量行为,如下所示:
>>> 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
这只是说如果你有一个名为 "vector2"
的记录,那么它应该被呈现为 Python 类Vector2
,如果你有一个"vector2"
记录数组(任何深度),它们应该集体呈现为一个ArrayOfVector2
类。
>>> 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...'>
这使您可以使用已定义的方法和属性:
>>> vectors[0, 0].norm
1.4866068747318506
>>> vectors.norm
<Array [[1.49, 2.97, 4.46], ... [5.95, 7.43]] type='3 * var * float64'>
而且,这正是https://github.com/scikit-hep/vector项目正在做的事情。密切关注它!