Search code examples
pythonawkward-array

while loop equivalent function/snippet in awkward array

I have a function which I would like to convert so that i can use with awkward array 1.

Function following which used to work for float but not for awkward arrays for the known reasons.

def Phi_mpi_pi(x):
    kPI=(3.14159265)
    kTWOPI = 2 * kPI

    #while ((x.any() >= kPI).any()): x = x - kTWOPI;                                                                                                                                                        
    #while ((x.any() < -kPI).any()): x = x + kTWOPI;                                                                                                                                                        
    while ((x >= kPI)): x = x - kTWOPI;
    while ((x < -kPI)): x = x + kTWOPI;
    return x;

I tried to convert it into numpy/awkward compatible form and new function look like

def Phi_mpi_pi(x):
    kPI=numpy.array(3.14159265)
    kPI = kPI.repeat(len(x))
    kTWOPI = 2 * kPI
    while ((x >= kPI)): x = x - kTWOPI;
    while ((x < -kPI)): x = x + kTWOPI;
    return x;

This function remains stuck in the while loop forever, I couldn't find a way to debug it.

Task of the function is to keep the values in an awkward array between +- kPI but this logic does not give the desired results.

e.g.

x=ak.Array([[0.7999999999999998, 1.0, -1.3], [], [-1.4], [-1.8000000000000003, -6.1000000000000005, -1.6000000000000005], [-4.6]])

However ((x < -kPI)) this give desired output.

>>> ak.to_list(x <= -kPI)
    [[False, False, False], [], [False], [False, True, False], [True]]

but not the function

the desired output should be b/w +- kPI based on the logic of while loop, is there something straightforward or suggestion which can be used?

Solution

  • Okay, got it. You want to adjust every single scalar value in x (which are all angles) to be between -π and π.

    You can do it like this:

    def Phi_mpi_pi(x):
    y = numpy.add(x, numpy.pi)
    y = numpy.mod(y, 2*numpy.pi)
    y = numpy.subtract(y, numpy.pi)
    return y

    Or, more terse and much less readable:

    def Phi_mpi_pi(x):
    return numpy.subtract(numpy.mod(numpy.add(x, numpy.pi), 2*numpy.pi), numpy.pi)

    What it does is this:

    1. Add π to all angles (so they point in the opposite direction).
    2. Take modulo 2π on all angles so they are all from 0 to 2π (2π not included).
    3. Subtract π from all the angles again (so they point in the right direction again). Now they are all from -π to +π (+π not included).

    Test:

    x = ak.Array([[0.3, 3.1, numpy.pi, -numpy.pi, -4 * numpy.pi,
               200 * numpy.pi, 2 * numpy.pi, -400 * numpy.pi], 
              [], [-1.4], [-1.8, -6, -1.6], [-4.6]])
y = Phi_mpi_pi(x)
print("Type of result:", type(y))
print("Result =", y)

# Check that the resulting range of each value is correct.
range_is_correct = (ak.all(y >= -numpy.pi) and ak.all(y < numpy.pi))

# Calculate the factors of the 2π adjustments.
factors = (x - y) / (2 * numpy.pi)
print("2π factors=", factors)

# Test that all factors of the 2π adjustmenst are approximate integers.
adjustments_are_correct = ak.all(numpy.abs(numpy.mod(factors, 1)) < 1e-15)

# Test that all values are correct, given the test input.
print("Result is correct:", range_is_correct and adjustments_are_correct)

    gives this output:

    Type of result: <class 'awkward1.highlevel.Array'>
Result = [[0.3, 3.1, -3.14, -3.14, 0, 1.78e-14, 0, 4.62e-14, ... [-1.8, 0.283, -1.6], [1.68]]
2π factors= [[2.65e-17, 7.07e-17, 1, 0, -2, 100, 1, -200], [], [0], [0, -1, 0], [-1]]
Result is correct: True

    which proves that the operation was correctly performed with the specifically used test data.

    Was that what you wanted?