def stream_function(step):
if step.geom.twod_yz:<EOL><INDENT>x_coord = step.geom.y_coord<EOL>v_x = step.fields['<STR_LIT>'][<NUM_LIT:0>, :, :, <NUM_LIT:0>]<EOL>v_z = step.fields['<STR_LIT>'][<NUM_LIT:0>, :, :, <NUM_LIT:0>]<EOL>shape = (<NUM_LIT:1>, v_x.shape[<NUM_LIT:0>], v_x.shape[<NUM_LIT:1>], <NUM_LIT:1>)<EOL><DEDENT>elif step.geom.twod_xz and step.geom.cartesian:<EOL><INDENT>x_coord = step.geom.x_coord<EOL>v_x = step.fields['<STR_LIT>'][:, <NUM_LIT:0>, :, <NUM_LIT:0>]<EOL>v_z = step.fields['<STR_LIT>'][:, <NUM_LIT:0>, :, <NUM_LIT:0>]<EOL>shape = (v_x.shape[<NUM_LIT:0>], <NUM_LIT:1>, v_x.shape[<NUM_LIT:1>], <NUM_LIT:1>)<EOL><DEDENT>else:<EOL><INDENT>raise NotAvailableError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>psi = np.zeros_like(v_x)<EOL>if step.geom.spherical:  <EOL><INDENT>r_nc = step.geom.r_coord + step.geom.rcmb  <EOL>r_pc = step.geom.r_mesh[<NUM_LIT:0>, <NUM_LIT:0>, :]  <EOL>r_nw = r_edges(step)[<NUM_LIT:0>][:<NUM_LIT:2>]  <EOL>vz0 = ((r_nw[<NUM_LIT:1>] - r_nc[<NUM_LIT:0>]) * v_z[:, <NUM_LIT:0>] +<EOL>(r_nc[<NUM_LIT:0>] - r_nw[<NUM_LIT:0>]) * v_z[:, <NUM_LIT:1>]) / (r_nw[<NUM_LIT:1>] - r_nw[<NUM_LIT:0>])<EOL>psi[<NUM_LIT:1>:, <NUM_LIT:0>] = -integrate.cumtrapz(r_pc[<NUM_LIT:0>]**<NUM_LIT:2> * vz0, x=x_coord)<EOL>vxc = (v_x + np.roll(v_x, -<NUM_LIT:1>, axis=<NUM_LIT:0>)) / <NUM_LIT:2><EOL>for i_x in range(len(x_coord)):<EOL><INDENT>psi[i_x, <NUM_LIT:1>:] = psi[i_x, <NUM_LIT:0>] +integrate.cumtrapz(r_pc * vxc[i_x], x=r_nc)<EOL><DEDENT><DEDENT>else:  <EOL><INDENT>z_nc = step.geom.z_coord<EOL>z_nw = r_edges(step)[<NUM_LIT:0>][:<NUM_LIT:2>]<EOL>vz0 = ((z_nw[<NUM_LIT:1>] - z_nc[<NUM_LIT:0>]) * v_z[:, <NUM_LIT:0>] +<EOL>(z_nc[<NUM_LIT:0>] - z_nw[<NUM_LIT:0>]) * v_z[:, <NUM_LIT:1>]) / (z_nw[<NUM_LIT:1>] - z_nw[<NUM_LIT:0>])<EOL>psi[<NUM_LIT:1>:, <NUM_LIT:0>] = -integrate.cumtrapz(vz0, x=x_coord)<EOL>vxc = (v_x + np.roll(v_x, -<NUM_LIT:1>, axis=<NUM_LIT:0>)) / <NUM_LIT:2><EOL>for i_x in range(len(x_coord)):<EOL><INDENT>psi[i_x, <NUM_LIT:1>:] = psi[i_x, <NUM_LIT:0>] +integrate.cumtrapz(vxc[i_x], x=z_nc)<EOL><DEDENT><DEDENT>if step.geom.twod_xz:<EOL><INDENT>psi = - psi<EOL><DEDENT>psi = np.reshape(psi, shape)<EOL>return psi<EOL>
Stream function.

    Args:
        step (:class:`~stagpy.stagyydata._Step`): a step of a StagyyData
            instance.
    Returns:
        :class:`numpy.array`: the stream function field, with four dimensions:
        x-direction, y-direction, z-direction and block.