Why not write the experimental code in the notebook, then write a function containing that code in the notebook and try it out, and only then move your function to a py file, if at all?