code: def plot_joint_sfs(s, ax=None, imshow_kwargs=None):
      
  
      import matplotlib.pyplot as plt
      from matplotlib.colors import LogNorm
  
      # check inputs
      s = asarray_ndim(s, 2)
  
      # setup axes
      if ax is None:
          w = plt.rcParams['figure.figsize'][0]
          fig, ax = plt.subplots(figsize=(w, w))
  
      # set plotting defaults
      if imshow_kwargs is None:
          imshow_kwargs = dict()
      imshow_kwargs.setdefault('cmap', 'jet')
      imshow_kwargs.setdefault('interpolation', 'none')
      imshow_kwargs.setdefault('aspect', 'auto')
      imshow_kwargs.setdefault('norm', LogNorm())
  
      # plot data
      ax.imshow(s.T, **imshow_kwargs)
  
      # tidy
      ax.invert_yaxis()
      ax.set_xlabel('derived allele count (population 1)')
      ax.set_ylabel('derived allele count (population 2)')
  
      return ax
docs: Plot a joint site frequency spectrum.
  
      Parameters
      ----------
      s : array_like, int, shape (n_chromosomes_pop1, n_chromosomes_pop2)
          Joint site frequency spectrum.
      ax : axes, optional
          Axes on which to draw. If not provided, a new figure will be created.
      imshow_kwargs : dict-like
          Additional keyword arguments, passed through to ax.imshow().
  
      Returns
      -------
      ax : axes
          The axes on which the plot was drawn.