How to clean the indexes, and ideally not create an additional array

Question

So I answered this question on SO and probably did someone's homework along the way. In the original question the OP has the answer variable and wants it split to real and imaginary values with the original coefficient (q in the original post) as an array for each q from 1 to 10, so you get an array of arrays looking like:

[[1,-0.39002422, 1.84237253, -0.39002422, -1.84237253, -0.10997578, 0.51949688,-0.10997578, -0.51949688j], 
[2, ...]
...
[10, ...]]

There are two things I'm unsatisfied with in the answer I gave that are the reason I'm asking for a CR.

1. The indexing: Isn't particularly readable, and I don't know how to make it more readable. For example [q-1, 2*_+1]. Would it have been better to create a variable j = q-1 because q-1 sees frequent use? Isn't that just moving the ugly to a different place?
2. The idea for the solution itself. The way I see it the original problem can be solved by splitting the answer array to a new array, but that feels somehow wrong. I can see several ways to create the new array, but I don't see a way to do it without creating the new array. Am I wrong and this is an OK solution or not, and why?

Thanks.

import numpy as np
export = np.empty([10,9])
for q in range(1,11):
    a0=1
    a1=3*q^2
    a2=4*q 
    a3=np.sqrt(q) 
    a4=q
    coeff=[a0, a1, a2, a3, a4]
    answer = np.roots(coeff)
    export[q-1, 0] = q
    for _ in range(len(answer)):
        export[q-1, 2*_+1] = answer[_].real
        export[q-1, 2*(_+1)] = answer[_].imag
with open("/tmp/latest_foo.csv", "ab") as output_file:
    np.savetxt(output_file, export, delimiter=",")

Answer 1

Here is a review of the solution.

• ^ is xor in Python. It is not for computation of exponentials.

• When running code outside a method / class, it is a good practice to put the code inside a main guard. See here for more explanation.

   if __name__ == "__main__":
      ...
  

  When you are providing quick answers on a forum, the guard might not always be necessary. However, if you are writing reusable code for yourself, it is better to adopt the practice.

• By convention, _ is used to represent a don't-care variable, which won't be used after its value is assigned, e.g. v = [[] for _ in range(4)]. It is undesirable to refer to it in an expression like answer[_].

• The output is a pure text file. It is unnecessary to open it in binary mode (b). The original question does not suggest an append mode (a) either.

• The q-1 index would no longer work if q is changed to a different group of values. Therefore, it is better to use enumerate in this case:

  for i, q in enumerate(RANGE_Q):
     ...
     export[i, 0] = q
     ...
  

  A better approach is to use zip:

  for q, output_row in zip(RANGE_Q, output_arr):
      ...
      output_row[0] = q
      ...
  

• The assignments to the output numpy array can be improved, as shown in my solution below.

Here is my solution.

import numpy as np

if __name__ == "__main__":
    # Define constants
    RANGE_Q = range(1, 11)               # Range of q values
    POLY_DEG = 4                         # Degree of polynomial
    OUTPUT_PATH = "/tmp/latest_foo.csv"  # Output csv path

    # Compute roots
    roots_arr = np.array([np.roots([1, 3*q*q, 4*q, np.sqrt(q), q]) for q in RANGE_Q])
    
    # Construct output array and assign values
    output_arr = np.empty(shape=(len(RANGE_Q), POLY_DEG * 2 + 1))
    output_arr[:, 0] = RANGE_Q
    output_arr[:, 1::2] = roots_arr.real
    output_arr[:, 2::2] = roots_arr.imag

    # Save results to output file
    np.savetxt(OUTPUT_PATH, output_arr, fmt="%.4g", delimiter=',')

Answer 2

A more pythonic method for that last for loop would be to use nested list comprehensions. It's likely faster, as well:

[item for sublist in [[x.real, x.imag] for x in answer] for item in sublist]

See this question for details about what's happening here, if you're not already familiar with list comprehensions.

In other words, instead of these four lines:

export[q-1, 0] = q
for _ in range(len(answer)):
      export[q-1, 2*_+1] = answer[_].real
      export[q-1, 2*(_+1)] = answer[_].imag

You could write this:

line = [q]
line.extend([item for sublist in [[x.real, x.imag] for x in answer] for item in sublist])
export[q-1] = line