I have 3 methods to do some simple calculations. 2 methods are returning nd-array and 1 method returns 1d-array. Later, I am creating a pandas dataframe based on the return output from the methods.

While I am creating pandas dataframe, I am also calculating std from the method's result. For nd-array I need to use axis=0 and axis=1 to calculate std but for the 1d-array, I can not use the axis properties.

That means I need to use if-else to calculate std for different returns from the methods. Below code is working fine

def main_fn(arr_1):
    all_result_summary = []
    for method in ["met_1", "met2", "met_3"]:
        results: ndarray = np.array(main_fn(list(arr_1), method))
        if method == "met_3":
            all_result_summary.append(
                    pd.DataFrame(
                        {
                            "Method": method,
                            "result": results.mean(),
                            "result_sd_ax_0": results.std(ddof=1),
                            "result_sd_ax_1": "NA",
                        },
                        index=[0],
                    )
            )
        else:
            all_result_summary.append(
                pd.DataFrame(
                    {
                        "Method": method,
                        "result": results.mean(),
                        "result_sd_ax_0": results.mean(axis=0).std(ddof=1),
                        "result_sd_ax_1": results.mean(axis=1).std(ddof=1),
                    },
                    index=[0],
                )
            )
        summary = pd.concat(all_result_summary, axis=0, ignore_index=True)
    return summary

However, I wanted to use a more pythonic way instead of reusing the whole code using if-else. Any other way to optimize the code?

CodePudding user response：

How about use x if cond else y, ternary-style syntax?

def main_fn(arr_1):
    all_result_summary = []
    for method in ["met_1", "met2", "met_3"]:
        results: ndarray = np.array(main_fn(list(arr_1), method))
        all_result_summary.append(
            pd.DataFrame(
                {
                    "Method": method,
                    "result": results.mean(),
                    "result_sd_ax_0": results.std(ddof=1) if method == "met_3" else results.mean(axis=0).std(ddof=1),
                    "result_sd_ax_1": "NA" if method == "met_3" else results.mean(axis=1).std(ddof=1),
                },
                index=[0],
            )
        )
        summary = pd.concat(all_result_summary, axis=0, ignore_index=True)
    return summary