pandas rename columns with method chaining

To rename columns in a chain

Use set_axis along axis=1:

df.set_axis(['foo', 'bar', 'baz'], axis=1)

If using groupby, pivot, melt, etc

Combine set_axis with pipe if the new columns depend on some earlier step in the chain. For example, to capitalize pivoted columns in a chain:

• We cannot directly chain set_axis:

  # does NOT work since df.columns are the original columns, not pivoted columns
  df.pivot(...).set_axis(df.columns.str.upper(), axis=1))
  

• But we can pipe the pivoted result into set_axis:

  # does work since we've piped the pivoted df
  df.pivot(...).pipe(lambda pivoted: pivoted.set_axis(pivoted.columns.str.upper(), axis=1)))
  

OP's example

Since OP has created a pivot_table and wants to conditionally collapse those pivoted MultiIndex, we pipe the pivot_table into the list comprehension:

(df.assign(date=pd.to_datetime(df.date))
   .assign(entries_per_ID=df.groupby('ID').ID.transform('size'))
   .pivot_table(index=['ID', 'date', 'code'],
                columns='location',
                values='entries_per_ID',
                aggfunc='max')
   .reset_index()
   .pipe(lambda piv: piv.set_axis(['_'.join(col).strip() if col[1] else col[0] for col in piv.columns],
                                  axis=1)))

#    ID        date   code  entries_per_ID_down  entries_per_ID_up
# 0   1  2021-10-12  False                  NaN                1.0
# 1   2  2021-10-15  False                  2.0                NaN
# 2   2  2021-10-16  False                  NaN                2.0
# 3   3  2021-10-01  False                  3.0                NaN
# 4   3  2021-10-10   True                  NaN                3.0
# 5   3  2021-10-19  False                  NaN                3.0