Search code examples
pythonnumpynumpy-ndarrayarray-broadcasting

Explicitly broadcasting a scalar


Is it possible to explicitly broadcast a scalar to fit an array somehow similar to

s[..., np.newaxis]

(So I want to add a dimension to s, even if it is only a scalar yet)

I am currently vectorizing a lot of functions, such that they work for a single datapoint or multiple at once. E.g. this function that either returns a single random normalized vector, or multiple at once:

import numpy as np
def randomu(N, M=None):
    """Returns one or M random normalized vectors of R^N."""
    v = np.random.normal(size=N if M is None else (M, N))
    return v / (np.linalg.norm(v) if M is None else np.linalg.norm(v, axis=-1)[:, np.newaxis])

One can see that this function is basically the single vector version and the multiple vectors version stitched together with two ternary condtionals. But if it would be possible to broadcast an additional axis to np.linalg.norm(v, axis=-1), no matter if it is a scalar or a vector, then a lot of my functions could be vectorized a lot cleaner.

To be precise: I want the function above to return an N-array for arguments N, None and to return a M,N-array for arguments N, M where M is some positive integer. (There is one dimension difference between M=Noneand M=1)

This would give this explanatory function a similar signature as most numpy functions as np.zeros, np.random.normal, .... Probably nearly all numpy functions that take an integer or tuple as shape argument.


Solution

  • It's not entirely clear which scalar you are worried about. Your v shape is either (N,) or (N,M). The only 'scalar' I see is the result of the norm, and that's a numpy.float64. It takes the [...,None] indexing:

    In [275]: np.linalg.norm(np.ones((3)), axis=-1)
    Out[275]: 1.7320508075688772
    In [276]: np.linalg.norm(np.ones((3)), axis=-1)[...,None]
    Out[276]: array([1.73205081])
    

    norm, like many "reduction" functions, takes a keepdims parameter, which ends up doing the same thing:

    In [277]: np.linalg.norm(np.ones((3)), axis=-1, keepdims=True)
    Out[277]: array([1.73205081])
    

    For a two array:

    In [278]: np.linalg.norm(np.ones((3,2)), axis=-1)
    Out[278]: array([1.41421356, 1.41421356, 1.41421356])    
    In [279]: np.linalg.norm(np.ones((3,2)), axis=-1)[...,None]
    Out[279]: 
    array([[1.41421356],
           [1.41421356],
           [1.41421356]])    
    In [280]: np.linalg.norm(np.ones((3,2)), axis=-1, keepdims=True)
    Out[280]: 
    array([[1.41421356],
           [1.41421356],
           [1.41421356]])
    

    These work with a v/..., since (3,) and (1,) broadcast, as do (3,2) and (3,1).

    def randomu(N, M=None):
        """Returns one or M random normalized vectors of R^N."""
        v = np.random.normal(size=N if M is None else (M, N))
        return v / np.linalg.norm(v, axis=-1, keepdims=True)
    
    In [284]: randomu(3)
    Out[284]: array([-0.92136225, -0.35801172,  0.15139094])
    
    In [285]: randomu(3,1)
    Out[285]: array([[-0.2261399 ,  0.16987621, -0.95916777]])
    
    In [286]: randomu(3,2)
    Out[286]: 
    array([[-0.73506825,  0.63749884, -0.23080272],
           [ 0.34294788,  0.93787639, -0.05267469]])