Add support for plotting isocontours of vector potential in snap.py

[soumide1102]

This PR enables snap.py to plot isocontours of A using the overlay_field routine in plot.py.

[Yurlungur]

Overall, looks good. A few comments below. And before I approve it, I want to test it myself.

[Yurlungur]

you may want to make zorder an optional parameter that can be passed in to the function, in case the user wants to layer them some different way.

[Yurlungur]

Why are you dividing hdr['N1'] by 2?

[soumide1102]

Yes, this is the line that I wanted to check with you.

Is A_phi[j,N1-1-i] calculating A in the -ve direction and A_phi[j,i+N1] in the +ve direction? So if I don't divide by 2, ultimately, in line 363, rcyl and z ends up having dimensions 256 X 256, and A_phi is 256 X 512. So I divided by 2 to make N1 represent cells in either the +ve/-ve direction only.

[Yurlungur]

Should be type=int and default=2

[soumide1102]

whoops! result of copy pasting the previous line to avoid typing, and then forgetting to update it completely. :)

[Yurlungur]

Looks like there's a row vs. column major ordering issue. N1 should be the the rows and N2 should be the columns, not the other way around. Not your fault. I'm sure it's my code that caused the problem there.

[Yurlungur]

Looks like there's a row vs. column major ordering issue. N1 should be the the rows and N2 should be the columns, not the other way around. Not your fault. I'm sure it's my code that caused the problem there.

Upon further analysis, the transpose isn't really the problem... whoever wrote the original code liked column-major order, I guess. But the transposes are all consistent. The real issue was that the indexing insanity for calculating A is only correct in 3d. In 2d it's wrong. I implement a fix by implementing a 2d and a 3d version below, and wrapping them. See:

def overlay_field_2d(ax, geom, dump, NLEV=20, linestyle='-', linewidth=1,
                  linecolor='k', zorder=2):
    from scipy.integrate import trapz
    hdr = dump['hdr']
    N1 = hdr['N1']; N2 = hdr['N2']
    x = geom['x'][:,:,0]
    y = geom['y'][:,:,0]
    z = geom['z'][:,:,0]
    rcyl = np.sqrt(x**2 + y**2)
    A_phi = np.zeros([N1,N2])
    gdet = geom['gdet']
    B1 = dump['B1'][:,:,0]
    B2 = dump['B2'][:,:,0]
    for i in range(N1):
        for j in range(N2):
            A_phi[i,j] = (trapz(gdet[:i,j]*B2[:i,j], dx=hdr['dx'][1]) - 
                          trapz(gdet[i,:j]*B1[i,:j], dx=hdr['dx'][2]))
    Apm = np.fabs(A_phi).max()
    if np.fabs(A_phi.min()) > A_phi.max():
        A_phi *= -1.
    levels = np.concatenate((np.linspace(-Apm,0,NLEV)[:-1], 
                             np.linspace(0,Apm,NLEV)[1:]))
    ax.contour(rcyl, z, A_phi, levels=levels, colors=linecolor, linestyles=linestyle,
               linewidths=linewidth, zorder=zorder)
    return

def overlay_field_3d(ax, geom, dump, NLEV=20, linestyle='-', linewidth=1,
                     linecolor='k', zorder=2):
  from scipy.integrate import trapz
  hdr = dump['hdr']
  N1 = hdr['N1']; N2 = hdr['N2']
  x = geom['x']
  y = geom['y']
  z = geom['z']
  x = flatten_xz(x, hdr, True).transpose()
  y = flatten_xz(x, hdr, True).transpose()
  z = flatten_xz(z, hdr, True).transpose()
  rcyl = np.sqrt(x**2 + y**2)
  A_phi = np.zeros([N2, 2*N1])
  gdet = geom['gdet'].transpose()
  B1 = dump['B1'].mean(axis=-1).transpose()
  B2 = dump['B2'].mean(axis=-1).transpose()
  for j in range(N2):
    for i in range(N1):
      A_phi[j,N1-1-i] = (trapz(gdet[j,:i]*B2[j,:i], dx=hdr['dx'][1]) - 
                         trapz(gdet[:j, i]*B1[:j, i], dx=hdr['dx'][2]))
      A_phi[j,i+N1] = (trapz(gdet[j,:i]*B2[j,:i], dx=hdr['dx'][1]) - 
                         trapz(gdet[:j, i]*B1[:j, i], dx=hdr['dx'][2]))
  A_phi -= (A_phi[N2//2-1,-1] + A_phi[N2//2,-1])/2.
  Apm = np.fabs(A_phi).max()
  if np.fabs(A_phi.min()) > A_phi.max():
    A_phi *= -1.
  #NLEV = 20
  levels = np.concatenate((np.linspace(-Apm,0,NLEV)[:-1], 
                           np.linspace(0,Apm,NLEV)[1:]))
  ax.contour(rcyl, z, A_phi, levels=levels, colors=linecolor, linestyles=linestyle,
             linewidths=linewidth, zorder=zorder)
  return

def overlay_field(ax, geom, dump, NLEV=20, linestyle='-', linewidth=1,
                  linecolor='k', zorder=2):
    if dump['hdr']['N3'] > 1. and dump['hdr']['stopx'][3] >= np.pi:
        return overlay_field_3d(ax,geom,dump,NLEV,linestyle,linewidth,linecolor,zorder)
    else:
        return overlay_field_2d(ax,geom,dump,NLEV,linestyle,linewidth,linecolor,zorder)

[Yurlungur]

I just noticed that args is referenced in make_snap now. This is going to break movie.py which calls this method. Can you pass these through as optional arguments to make_snap so that the args object isn't used in scope it might not be available?

[Yurlungur]

@soumide1102 what's the status of this MR?