Seb wrote: > Hello, > > I'm fairly new to Python, struggling to write in a more object-oriented, > functional style. I just wrote a function that takes two arrays > representing sine (y) and cosine (x) angle coordinates, and returns the > angle in degrees. I had initially written the function to take > array-like arguments x/y, but I'd like to generalize and take scalars as > well. However, the function has a subsetting operations, which don't > work with scalars: > > vmag = np.sqrt((x ** 2) + (y ** 2)) > ang = np.arctan2(y, x) > ang[ang < 0] = ang[ang < 0] + (2 * np.pi) # output range 0 - 2*pi > ang[vmag == 0] = 0 # when magnitude is 0 the angle is also 0 > ang[ang == 0] = 2 * np.pi # convention > > If I want to take scalars x/y, I naively thought about implementing an > if/else statement right before the subsetting operations. However, my > intuition tells me there must be a proper object-oriented solution to > this. Any tips appreciated.
Here's a "duck-typed" solution: import numpy as np def _f(x, y): vmag = np.sqrt((x ** 2) + (y ** 2)) ang = np.arctan2(y, x) ang[ang < 0] = ang[ang < 0] + (2 * np.pi) # output range 0 - 2*pi ang[vmag == 0] = 0 # when magnitude is 0 the angle is also 0 ang[ang == 0] = 2 * np.pi # convention return ang def f(x, y): try: return _f(x, y) except IndexError: return _f(np.array([x]), y)[0] However, this is quite inefficient for scalars. Personally I'd implement two versions, one for scalars and one for vectors. This may be a little less convenient, but most of the time the client code has to deal either with scalars or vectors, not both. -- https://mail.python.org/mailman/listinfo/python-list