Save points cloud in PLOT3D format

Question

I have the following function to save points cloud in PLOT3D format:

def write(filename, mesh):
    """
    Write the mesh to filename in PLOT3D format.
    filename: is the name of the output file.
    mesh: is a numpy array of size (N,2) or (N, 3) where N is the number of points. mesh[i] is a numpy array coordinates of the i'th point.

    """
    points = mesh.points
    imax = np.unique(points[:, 0]).size
    jmax = np.unique(points[:, 1]).size

    _2d = False
    if points.shape[1] == 2:
        _2d = True
    elif np.unique(points[:, 2]).size == 1:
        _2d = True

    with open(filename, "w") as p3dfile:
        if not _2d:
            kmax = np.unique(points[:, 2]).size
            print(imax, jmax, kmax, file=p3dfile)
            for value in points.flatten(order="F"):
                print(value, file=p3dfile)
        else:
            print(imax, jmax, file=p3dfile)
            for value in points[:, 0:2].flatten(order="F"):
                print(value, file=p3dfile)

For large number of points this function is very slow.

I appreciate any suggestion to improve the code above.

Answer 1

Optimizations/improvements:

Namings

• write is too generic name for specific function. The one of better names is write_points
• _2d. Even if a variable is marked as "protected" _..., starting its name with a digit is a bad naming pattern. As it's a boolean flag a meaningful name would be is_2dim (is 2-dimensional space)

Since mesh argument is only used to access its inner property points it's better to pass the target points data structure directly (It's called Replace Query with Parameter technique)

In current approach, when reaching the condition if not _2d: the expression np.unique(points[:, 2]).size will be calculated repeatedly, instead:
since points array is expected to be 2d array with either 2 or 3 columns we can collect unique counts across all columns at once (instead of declaring imax, jmax):

col_counts = [np.unique(row).size for row in points.T]

The whole conditional:

_2d = False
if points.shape[1] == 2:
    _2d = True
elif np.unique(points[:, 2]).size == 1:
    _2d = True

is now replaced with a single statement:

is_2dim = points.shape[1] == 2 or col_counts[2] == 1

Calling print function at each iteration to write the data into file:

for value in points.flatten(order="F"):
    print(value, file=p3dfile)

is definitely less efficient and performant than calling writelines() on generator expression at once.

---

The final optimized function:

def write_points(filename, points):
    """
    Write the mesh points to filename in PLOT3D format.
    filename: is the name of the output file.
    mesh: is a numpy array of size (N,2) or (N, 3) where N is the number of points. mesh[i] is a numpy array coordinates of the i'th point.

    """
    col_counts = [np.unique(row).size for row in points.T]
    is_2dim = points.shape[1] == 2 or col_counts[2] == 1

    with open(filename, "w") as p3dfile:
        pd3file.write(' '.join(map(str, col_counts)))
        pd3file.writelines(f'{num}\n' for num in points.flatten(order="F"))