I'm fairly new to NumPy, so it's quite possible that I'm missing something fundamental. Don't hesitate to ask "stupid" questions about "basic" things!
I'm trying to write some functions that manipulate vectors. I'd like them to work on single vectors, as well as on arrays of vectors, like most of NumPy's ufuncs:
import math
import numpy
def func(scalar, x, vector):
# arbitrary function
# I'm NOT looking to replace this with numpy.magic_sum_multiply()
# I'm trying to understand broadcasting & dtypes
return scalar * x + vector
print(func(
scalar=numpy.array(2),
x=numpy.array([1, 0, 0]),
vector=numpy.array([1, 0, 0]),
))
# => [3 0 0], as expected
print(func(
scalar=numpy.array(2),
x=numpy.array([1, 0, 0]),
vector=numpy.array([[1, 0, 0], [0, 1, 0]]),
))
# => [[3 0 0], [2 1 0]], as expected. x & scalar are broadcasted out to match the multiple vectors
However, when trying to use multiple scalars, things go wrong:
print(func(
scalar=numpy.array([1, 2]),
x=numpy.array([1, 0, 0]),
vector=numpy.array([1, 0, 0]),
))
# => ValueError: operands could not be broadcast together with shapes (2,) (3,)
# expected: [[2 0 0], [3 0 0]]
I'm not entirely surprised be this. After all, NumPy has no idea that I'm working with vectors that are an single entity, and not an arbitrary dimension.
I can solve this ad-hoc with some expand_dims()
and/or squeeze()
to add/remove axes, but that feels hacky...
So I figured that, since I'm working with vectors that are a single "entity", dtypes may be what I'm looking for:
vector_dtype = numpy.dtype([
('x', numpy.float64),
('y', numpy.float64),
('z', numpy.float64),
])
_ = numpy.array([(1, 0, 0), (0, 1, 0)], dtype=vector_dtype)
print(_.shape) # => (2,), good, we indeed have 2 vectors!
_ = numpy.array((1, 0, 0, 7), dtype=vector_dtype)
# Good, basic checking that I'm staying in 3D
# => ValueError: could not assign tuple of length 4 to structure with 3 fields.
However, I seem to loose basic math capabilities:
print(2 * _)
# => TypeError: The DTypes <class 'numpy.dtype[void]'> and <class 'numpy.dtype[uint8]'> do not have a common DType. For example they cannot be stored in a single array unless the dtype is `object`.
So my main question is: How do I solve this?
- Is there some numpy.magic_broadcast_that_understands_what_I_mean() function?
- Can I define math-operators (such as addition, ...) on the vector-dtype?