docs: def stream_function(step):
      """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.
      """
      if step.geom.twod_yz:
          x_coord = step.geom.y_coord
          v_x = step.fields['v2'][0, :, :, 0]
          v_z = step.fields['v3'][0, :, :, 0]
          shape = (1, v_x.shape[0], v_x.shape[1], 1)
      elif step.geom.twod_xz and step.geom.cartesian:
          x_coord = step.geom.x_coord
          v_x = step.fields['v1'][:, 0, :, 0]
          v_z = step.fields['v3'][:, 0, :, 0]
          shape = (v_x.shape[0], 1, v_x.shape[1], 1)
      else:
          raise NotAvailableError('Stream function only implemented in '
                                  '2D cartesian and spherical annulus')
      psi = np.zeros_like(v_x)
      if step.geom.spherical:  # YZ annulus
          # positions
          r_nc = step.geom.r_coord + step.geom.rcmb  # numerical centers
          r_pc = step.geom.r_mesh[0, 0, :]  # physical centers
          r_nw = r_edges(step)[0][:2]  # numerical walls of first cell
          # vz at center of bottom cells
          vz0 = ((r_nw[1] - r_nc[0]) * v_z[:, 0] +
                 (r_nc[0] - r_nw[0]) * v_z[:, 1]) / (r_nw[1] - r_nw[0])
          psi[1:, 0] = -integrate.cumtrapz(r_pc[0]**2 * vz0, x=x_coord)
          # vx at center
          vxc = (v_x + np.roll(v_x, -1, axis=0)) / 2
          for i_x in range(len(x_coord)):
              psi[i_x, 1:] = psi[i_x, 0] + \
                  integrate.cumtrapz(r_pc * vxc[i_x], x=r_nc)
      else:  # assume cartesian geometry
          z_nc = step.geom.z_coord
          z_nw = r_edges(step)[0][:2]
          vz0 = ((z_nw[1] - z_nc[0]) * v_z[:, 0] +
                 (z_nc[0] - z_nw[0]) * v_z[:, 1]) / (z_nw[1] - z_nw[0])
          psi[1:, 0] = -integrate.cumtrapz(vz0, x=x_coord)
          # vx at center
          vxc = (v_x + np.roll(v_x, -1, axis=0)) / 2
          for i_x in range(len(x_coord)):
              psi[i_x, 1:] = psi[i_x, 0] + \
                  integrate.cumtrapz(vxc[i_x], x=z_nc)
      if step.geom.twod_xz:
          psi = - psi
      psi = np.reshape(psi, shape)
      return psi