Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ENH: broadcast_arrays() with option to not repeat arrays #28098

Open
34j opened this issue Jan 4, 2025 · 4 comments
Open

ENH: broadcast_arrays() with option to not repeat arrays #28098

34j opened this issue Jan 4, 2025 · 4 comments
Labels
57 - Close? Issues which may be closable unless discussion continued

Comments

@34j
Copy link

34j commented Jan 4, 2025

Proposed new feature or change:

Consider a function $f(r, \theta) = [r \cos \theta, r \sin \theta]$ and when r.shape == (3, 1) and theta.shape == (4,) thus the shape of the result is supposed to be (2, 3, 4) (or (3, 4, 2)) by broadcasting.

import numpy as np


def ferror(r, theta): # (3, 1) and (4,)
	cossintheta = np.stack([np.cos(theta), np.sin(theta)]) # (2, 4) (or (4, 2) if axis=-1)
	return r * cossintheta # (3, 1) and (2, 4) (or (4, 2)) -> error

def fslow(r, theta): # (3, 1) and (4,)
	r, theta = np.broadcast_arrays(r, theta) # (3, 4) and (3, 4)
	cossintheta = np.stack([np.cos(theta), np.sin(theta)]) # (2, 3, 4)
	return r * cossintheta

def ffast(r, theta): # (3, 1) and (4,)
	r, theta = np.broadcast_arrays(r, theta, no_repeat=True) # (3, 1) and (1, 4)
	cossintheta = np.stack([np.cos(theta), np.sin(theta)]) # (2, 1, 4)
	return r * cossintheta

Where broadcast_arrays(*arrays, no_repeat=True) does

  • check if the arrays are broadcastable
  • for each array, add np.newaxis to the head until the dimension gets max([a.ndim for a in arrays])

Apparently the latter function ffast is less computationally demanding than fslow. Hence, it would be useful to have such an option.

import numpy as np

def broadcast_without_repeating(*arrays):
    np.broadcast_shapes([a.shape for a in arrays])    
    max_dim = max(a.ndim for a in arrays)
    return tuple(array[(None,) * (max_dim - array.ndim), ...] for array in arrays)
@mhvk
Copy link
Contributor

mhvk commented Jan 4, 2025

This function you seem to want is not really broadcasting, but more changing the number of dimensions, equivalent to

def same_ndim(*arrays):
    max_dim = max([a.ndim for a in arrays])
    return tuple(np.array(a, ndmin=max_dim, copy=False, subok=True) for a in arrays)

This is simple enough that I'm not sure it is worth adding (note that np.broadcast_arrays is not so simple, as it also is setting strides).

Another function that may help in cases like these is np.expand_dims or simply inserting extra axes with [:, np.newaxis], etc. In particular, if you use np.stack(..., axis=-1), then you know you want an extra trailing dimension on r as well, i.e., r = r[..., np.newaxis], after which you can just let regular broadcasting take care.

p.s. Your specific problem is probably just a simplest version of what you ran into, but it could be solved by using np.multiply.outer, as in,

def fouter(r, theta): # (3, 1) and (4,)
	cossintheta = np.stack([np.cos(theta), np.sin(theta)]) # (2, 4) (or (4, 2) if axis=-1)
	return np.multply.outer(r, cossintheta)  # shape 3, 2, 4 or 2,4,3 if order switched.

@34j
Copy link
Author

34j commented Jan 5, 2025

def same_ndim(*arrays):
max_dim = max([a.ndim for a in arrays])
return tuple(np.array(a, ndmin=max_dim, copy=False, subok=True) for a in arrays)

This function is certainly acceptable, but I would like to clarify what it would be better to check if the shapes are broadcastable in advance to prevent unnecessary calculations.

Furthermore, such a function can be used for vectorized arguments as well without increasing computational cost; Consider a function $g(r, (\theta, \phi)) = (r \sin \theta \cos \phi, r \sin \theta \sin \phi, r \cos \theta)$ when r.shape == (3, 1), thetaandphi.shape == (4, 2). (Suppose that for some reason θ and φ are stacked on axis -1. Not a very good example.)

def gfast(r, thetaandphi): # (3, 1) and (4, 2)
	"""r.shape and theta.shape[:-1] should be broadcastable and theta.shape[-1] should be 2"""
	if thetaandphi.shape[-1] != 2:
		raise ValueError()
	r, thetaandphi = np.broadcast_arrays(r[..., None], thetaandphi, no_repeat=True) # (3, 1, 1) and (1, 4, 2)
	r = r[..., 0] # (3, 1) # Thanks to no_repeat, this is not repeated AS WELL. 
	# Therefore, by using the above function, we are killing two birds with one stone.
	sintheta = np.sin(thetaandphi[..., 0])
	direction = np.stack([sintheta * np.cos(thetaandphi[..., 1]), sintheta * np.sin(thetaandphi[..., 1]), np.cos(thetaandphi[..., 0])])
	return r * direction

Without using such a function, broadcast_to() must be used, which increases the number of lines.

At any rate, if such a function were available in the standard API, users might be able to create functions with batch support (and/or vectorized arguments) more comfortably. (Although it would be difficult to do so)

@seberg
Copy link
Member

seberg commented Jan 6, 2025

I tend to agree with Marten that this doesn't seem all that common, and getting to the identical dimension (appending them) is easy enough.
FWIW, the last example:

r, thetaandphi = np.broadcast_arrays(r[..., None], thetaandphi, no_repeat=True) # (3, 1, 1) and (1, 4, 2)

with inputs # (3, 1) and (4, 2) doesn't look like broadcasting at all.

You might want something like a, b = np.ix_(a, b) also, in case you want the "outer" semantics (i.e. explicitly not broadcasting).

The thing that makes it very hard for me to even give a hint, is that I am not sure how to document what the function returns? The broadcast of r and theta with a new dimension of length 2? theta plus a new dim of length 2 then the outer with r (adding it's dimensions`)?

but I would like to clarify what it would be better to check if the shapes are broadcastable in advance to prevent unnecessary calculations.

While that may make sense, doing so seems unnecessary in most use-cases. Error cases like this should only be taken on programming issues (normally). If the user really can run into shape mismatch errors, it might be that they should check ahead of time (rather than this function).

@seberg seberg added the 57 - Close? Issues which may be closable unless discussion continued label Jan 6, 2025
@34j
Copy link
Author

34j commented Jan 6, 2025

with inputs # (3, 1) and (4, 2)

the inputs are actually (3, 1, 1) and (4, 2) because [..., None] is applied.

You might want something like a, b = np.ix_(a, b)

In my understanding np.ix_ takes only 1d arguments, and I am not sure what this has to do with this issue.

I am not sure how to document what the function returns?

It is quite simple, for the last example the return shape is (3,) + np.broadcast_shapes(r.shape, thetaandphi[..., 0]) (but I think it should have been np.broadcast_shapes(r.shape, thetaandphi[..., 0]) + (3,)).

While that may make sense, doing so seems unnecessary in most use-cases. Error cases like this should only be taken on programming issues (normally). If the user really can run into shape mismatch errors, it might be that they should check ahead of time (rather than this function).

  1. For more complex and slow functions, it is better to check the shapes at the beginning of the function. And it should be in accordance with the basic principles of the programming
  2. From personal experience, most of the errors I encounter when writing Numpy code are shape mismatches caused by mis-ordering of dimensions.
  3. This facilitates vectorization of calculation. In my experience, Numpy code created by inexperienced users often does not support vectorized input and uses dirty and slow for loops. If this is implemented in the API it would make the situation a bit better.
  4. I believe it is not a good idea to write many lines of checking code for each function or maintain a unique utility function for each library to check.

I strongly hope that this could be incorporated into the standard API. If not, is there a more appropriate place to implement this?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
57 - Close? Issues which may be closable unless discussion continued
Projects
None yet
Development

No branches or pull requests

3 participants