1 tplarson 1.1 import numpy as np
  2              import os
  3              import h5py
  4              import matplotlib.pyplot as plt
  5              #from matplotlib import animation

  6 tplarson 1.2 from matplotlib import cm
  7              from matplotlib.colors import ListedColormap

  8 tplarson 1.1 

  9 tplarson 1.2 modeldir='../mods'

 10 tplarson 1.1 modeln=np.array([])
 11              modell=np.array([])
 12              modelnu=np.array([])
 13              rmesh=np.array([])
 14              rho=np.array([])
 15              R=0.0
 16              
 17              def loadmodel():
 18              
 19              #  in_dir = '/home/tplarson/solar/mods'

 20 tplarson 1.2   in_dir = modeldir

 21 tplarson 1.1   fname = 'mods_p3_eigfcn.h5'
 22                hf = h5py.File(os.path.join(in_dir, fname))
 23              
 24              # load r, rho and R from "model"
 25                global rmesh, rho, R
 26                rmesh, rho, R = [hf['model'][k].value for k in ['r', 'rho', 'R']]
 27                rmesh /= R
 28              
 29              # load l and n from "modes"
 30                global modeln, modell, modelnu
 31                modeln, modell, modelnu = [hf['modes'][k].value for k in ['n', 'l', 'nu']]
 32              
 33              
 34              def getradial(l,n):
 35              
 36              #  in_dir = '/home/tplarson/solar/mods'

 37 tplarson 1.2   in_dir = modeldir

 38 tplarson 1.1   fname = 'mods_p3_eigfcn.h5'
 39                hf = h5py.File(os.path.join(in_dir, fname))
 40              
 41              # open y1 and y2 data sets from "modes" (without reading the data)
 42                y1ds = hf['modes/y1'] 
 43                y2ds = hf['modes/y2']
 44              
 45                idx = ((modeln == n) & (modell == l))
 46                y1, y2 = y1ds[idx,:], y2ds[idx,:]
 47              
 48              # compute xi_r and xi_h
 49                L2 = l * (l + 1)
 50                xir = (y1 * R).flatten()
 51                if (l > 0):
 52                  xih = (y2 * R / L2).flatten()
 53                else:
 54                  xih = 0.0*xir
 55              
 56                return (xir, xih)
 57              
 58              

 59 tplarson 1.2 def writesurfacemodel(lmin=0, lmax=300, rsurf=1.0):

 60 tplarson 1.1 

 61 tplarson 1.2   outfile='model.surface.modes'

 62 tplarson 1.1   file=open(outfile,'w')

 63 tplarson 1.2   in_dir = modeldir

 64 tplarson 1.1   fname = 'mods_p3_eigfcn.h5'
 65                hf = h5py.File(os.path.join(in_dir, fname))
 66              
 67              # load r, rho and R from "model"
 68                rmesh, c, R = [hf['model'][k].value for k in ['r', 'c', 'R']]
 69                indsurf = np.abs(rmesh-rsurf*R).argmin()
 70              
 71              # load l and n from "modes"
 72                modeln, modell, modelnu, modelsig2, modelE = [hf['modes'][k].value for k in ['n', 'l', 'nu', 'sigma2', 'E']]
 73              
 74                outfmt='{:d} {:d} {:f} {:f} {:f} {:f} {:f} {:f} {:f}'

 75 tplarson 1.2   l=lmin

 76 tplarson 1.1   while (l <= lmax):
 77                  ind = (modell == l)
 78                  nlist = modeln[ind]
 79                  for n in nlist:
 80                    xir, xih = getradial(l,n)
 81              #      rc=float(xir[indsurf])
 82              #      hc=float(xih[indsurf])
 83                    rc=np.interp(rsurf*R,rmesh,xir)
 84                    hc=np.interp(rsurf*R,rmesh,xih)
 85                    ind2 = (modell == l) & (modeln == n)
 86                    nu = float(modelnu[ind2])
 87                    erg = float(modelE[ind2])*1e6
 88                    sig2 = float(modelsig2[ind2])
 89                    if (l > 0):
 90                      L=np.sqrt(l*(l+1))
 91              #        indturn=np.abs(rmesh/c - L/(2*np.pi*nu)).argmin()
 92              #        rturn=float(rmesh[indturn]/R)
 93                      rturn=np.interp(L/(2*np.pi*nu),rmesh/c,rmesh)/R
 94                    else:
 95                      rturn=0.0
 96                    amp=nu*np.sqrt(np.square(rc)+l*(l+1)*np.square(hc))
 97 tplarson 1.1       nu*=1e6
 98                    amp=float(amp)/30e6
 99              #      outstr=f'{l:d} {n:d} {nu:f} {amp:f} {sig2:f} {rc:f} {hc:f} {rturn:f} {erg:f}'
100                    outstr=outfmt.format(int(l),int(n),nu,amp,sig2,rc,hc,rturn,erg)
101                    file.write('%s\n' % (outstr))
102                  l+=1
103                file.close()
104              
105              
106              
107              def image2sphere(xpixels=1000, ypixels=1000, distobs=220.0, \
108                               bangle=-30.0, pangle=0.0):
109              
110                p = pangle*np.pi/180
111                b0 = bangle*np.pi/180
112                #distobs = 220.0  # solar radii
113                #xpixels = 1000
114                #ypixels = 1000 
115                x0 = xpixels/2-0.5
116                y0 = ypixels/2-0.5
117                imscale = 1.97784*1024/xpixels
118 tplarson 1.1   scale = imscale/(180*60*60/np.pi)
119                rsun=np.tan(np.arcsin(1/distobs))/scale
120              
121              # Sun is sitting at the center of the main coordinate system and has radius 1.
122              # Observer is at robs=(xobs,yobs,zobs) moving with a velocity vobs.
123              # Start by setting robs from distobs and b0
124              
125                robs_x = distobs * np.cos(b0)
126                robs_y = 0.0
127                robs_z = distobs * np.sin(b0)
128              
129              # image coordinates
130                xx = np.linspace(0, xpixels-1, xpixels)
131                yy = np.linspace(0, ypixels-1, ypixels)
132                xx, yy = np.meshgrid(xx, yy)
133              
134                x2 = scale*(xx - x0)
135                y2 = scale*(yy - y0)
136              # Rotate by the P-angle. New coordinate system has the y-axis pointing
137              # towards the solar north pole.
138                x1 =  x2*np.cos(p) + y2*np.sin(p)
139 tplarson 1.1   y1 = -x2*np.sin(p) + y2*np.cos(p)
140              
141              # Now transform to put the coordinates into the solar coordinate system.
142              # First find the directions (vecx and vecy) the x2 and y2 coordinate
143              # axis correspond to. vecsun points towards the Sun. Note that the
144              # (x2,y2,Sun) system is left handed. These vectors are unit vectors.
145              
146                vecx_x=0.0
147                vecx_y=1.0
148                vecx_z=0.0
149                vecy_x=-np.sin(b0)
150                vecy_y=0.0
151                vecy_z=np.cos(b0)
152                vecsun_x=-np.cos(b0)
153                vecsun_y=0.0
154                vecsun_z=-np.sin(b0)
155              
156              # Now the proper direction can be found. These are not unit vectors.
157                x = vecx_x*x1 + vecy_x*y1 + vecsun_x
158                y = vecx_y*x1 + vecy_y*y1 + vecsun_y
159                z = vecx_z*x1 + vecy_z*y1 + vecsun_z
160 tplarson 1.1   qq = 1/np.sqrt(x*x + y*y + z*z)
161              # Make them unit vectors.
162                x*=qq
163                y*=qq
164                z*=qq
165              
166              # Now find intersection with the Sun.
167              # Solve quadratic equation |robs+q*[x1,y1,z1]|=1 for q
168              # a, b and c are terms in a*x^2+bx+c=0. a==1 since [x1,y1,z1] is unit vector.
169                c = robs_x*robs_x + robs_y*robs_y + robs_z*robs_z -1
170                b = 2*(x*robs_x+y*robs_y+z*robs_z)
171                d = b*b - 4*c
172                index = (d >= 0)
173                q = np.zeros([xpixels,ypixels])
174                xsun = np.zeros([xpixels,ypixels])
175                ysun = np.zeros([xpixels,ypixels])
176                zsun = np.zeros([xpixels,ypixels])
177                q[index]=(-b[index] - np.sqrt(d[index]))/2
178                xsun[index]=robs_x + x[index]*q[index]
179                ysun[index]=robs_y + y[index]*q[index]
180                zsun[index]=robs_z + z[index]*q[index]
181 tplarson 1.1 
182                phisun = np.arctan2(ysun,xsun)
183                thetasun = np.pi/2 - np.arcsin(zsun)
184              
185                ph=np.ma.array(phisun, mask=np.logical_not(index))
186                th=np.ma.array(thetasun, mask=np.logical_not(index))
187              
188                return (ph, th) 
189              
190              
191              def image2rtheta(xpixels=1000, ypixels=1000, distobs=220.0):
192              
193                imscale = 1.97784*1024/xpixels
194                scale = imscale/(180*60*60/np.pi)
195                rsun=np.tan(np.arcsin(1/distobs))/scale
196                x0 = xpixels/2-0.5
197                y0 = ypixels/2-0.5
198                xx = (np.linspace(0, xpixels-1, xpixels)-x0)/rsun
199                yy = (np.linspace(0, ypixels-1, ypixels)-y0)/rsun
200                xx, yy = np.meshgrid(xx, yy)
201                rr = np.sqrt(xx*xx+yy*yy)
202 tplarson 1.1   index = (rr <= 1.0)
203                r = np.ma.array(rr, mask=np.logical_not(index))

204 tplarson 1.2 #  lat = np.arctan(yy/np.abs(xx))
205                lat = np.arctan2(yy,np.abs(xx))

206 tplarson 1.1   theta = np.ma.array(np.pi/2 - lat, mask=np.logical_not(index))
207              
208                return (r, theta)
209              
210              
211              def setplm(l, m, x, plm, dplm): 
212              
213              # adapted from setplm.pro by Jesper Schou
214              # Set plm(*,l)=P_l^m (x) for l=m,lmax
215              # optionally sets dplm(*,l)={{dP_l^m} \over dx} (x) for l=m,lmax
216              # P_l^m 's are normalized to have \int_{-1}^1 (P_l^m (x))^2 dx = 1
217              # Works up to l \approx 1800, see ~/invnew/plm.f for details
218              
219                eps = 1.0e-12
220                x1 = np.maximum(np.minimum(x,1-eps),eps-1)

221 tplarson 1.2 #  x1 = x

222 tplarson 1.1   x2 = 1.0/(x1*x1-1.0)
223                x3 = x1*x2
224                c = np.sqrt((2*m+1)/2.0)
225                for i in range(1,m+1):
226                  c *= -np.sqrt(1.0-0.5/i)
227                plm[...,0] = c*(np.sqrt(1.0-x1*x1))**m
228                if (l > m):
229                  c = np.sqrt(2.0*m+3.0)
230                  plm[...,1] = c*x1*plm[...,0]
231                i = m+2
232                while (i <= l):
233                  c1 = np.sqrt((4.0*i*i-1.0)/(i*i-m*m))
234                  c2 = np.sqrt(((2.0*i+1.0)*(i+m-1.0)*(i-m-1.0))/((2.0*i-3.0)*(i*i-m*m)))
235                  plm[...,i-m] = c1*x1*plm[...,i-m-1] - c2*plm[...,i-m-2]
236                  i+=1
237              
238                dplm[...,0] = m*x3*plm[...,0]
239                i = m+1
240                while (i <= l):
241                  c = np.sqrt((2.0*i+1.0)*(i*i-m*m)/(2.0*i-1))
242                  dplm[...,i-m] = i*x3*plm[...,i-m] - c*x2*plm[...,i-m-1]
243 tplarson 1.1     i+=1
244              
245                return (plm[...,l-m],dplm[...,l-m])
246              
247              
248              lsave=0
249              msave=0
250              nsave=1
251              
252              def querylmn():
253              
254                global lsave, msave, nsave
255                lstr=catchc("Enter spherical harmonic degree (l): ",lsave)
256                if lstr == 'q':
257                  return (-1,-1,-1)
258                else:
259                  while True:
260                    try:
261                      lval=int(lstr)
262                      if lval < 0:
263                        print("Degree must be greater than or equal to zero. Try again.")
264 tplarson 1.1           lstr=catchc("Enter spherical harmonic degree (l): ",lsave)
265                        if lstr == 'q':
266                          return (-1,-1,-1)
267                        continue
268                      break
269                    except:
270                      print("Invalid number, try again.")
271                      lstr=catchc("Enter spherical harmonic degree (l): ", lsave)
272                      if lstr == 'q':
273                        return (-1,-1,-1)
274              
275                mstr=catchc("Enter azimuthal order (m): ",msave)
276                if mstr == 'q':
277                  return (-1,-1,-1)
278                else:
279                  while True:
280                    try:
281                      mval=int(mstr)
282                      if np.abs(mval) > lval:
283                        print("Azimuthal order must be in the range -l <= m <= l. Try again.")
284                        mstr=catchc("Enter azimuthal order (m): ",msave)
285 tplarson 1.1           if mstr == 'q':
286                          return (-1,-1,-1)
287                        continue
288                      break
289                    except:
290                      print("Invalid number, try again.")
291                      mstr=catchc("Enter azimuthal order (m): ",msave)
292                      if mstr == 'q':
293                        return (-1,-1,-1)
294              
295                nstr=catchc("Enter radial order (n): ",nsave)
296                if nstr == 'q':
297                  return (-1,-1,-1)
298                else:
299                  while True:
300                    try:
301                      nval=int(nstr)
302                      if nval < 0:
303                        print("Radial order must be greater than or equal to zero. Try again.")
304                        nstr=catchc("Enter radial order (n): ",nsave)
305                        if nstr == 'q':
306 tplarson 1.1             return (-1,-1,-1)
307                        continue
308                      break
309                    except:
310                      print("Invalid number, try again.")
311                      nstr=catchc("Enter radial order (n): ",nsave)
312                      if nstr == 'q':
313                        return (-1,-1,-1)
314              
315                lsave=lval
316                msave=mval
317                nsave=nval
318                return (lval,mval,nval)
319              
320              
321              varlist=['Vr', 'Vt', 'Vp', 'Vh', 'Vmag', 'Vsq']
322              colormap="seismic"
323              plotvar="Vr"
324              isave=0
325              
326              def queryplotparms():
327 tplarson 1.1 
328                global colormap, plotvar, isave
329                print("You may enter 'l' to list options.")
330                cmap=input("Enter name of colormap: ")
331                while True:
332                  if (cmap == 'q'):
333                    return cmap
334                  if (cmap == ''):
335                    print("Using saved value colormap = %s." % colormap)
336                    cmap=colormap
337                    break
338                  if (cmap == 'l'):
339                    printcmaps()
340                    print("Current colormap is %s." % colormap)
341                    cmap=input("Enter name of colormap: ")
342                    continue
343                  elif (cmap not in plt.colormaps()):
344                    print("That colormap not registered, try again.")
345                    cmap=input("Enter name of colormap: ")
346                    continue
347                  else:
348 tplarson 1.1       break
349                colormap=cmap
350                  
351                pvar=input("Enter variable to plot: ")
352                while True:
353                  if (pvar == 'q'):
354                    return pvar
355                  if (pvar == ''):
356                    print("Using saved value plotvar = %s." % plotvar)
357                    pvar=plotvar
358                    break
359                  if (pvar == 'l'):
360                    print("Available plotting variables are: ", end="")
361                    for i in varlist[:len(varlist)-1]:
362                      print(i, end=", ")
363                    print("and %s" % varlist[len(varlist)-1], end=".\n")
364                    print("Currently plotting %s." % plotvar)
365                    pvar=input("Enter variable to plot: ")
366                    continue
367                  elif (pvar not in varlist):
368                    print("That variable not available, try again.")
369 tplarson 1.1       pvar=input("Enter variable to plot: ")
370                    continue
371                  else:
372                    break
373                plotvar=pvar
374              
375                ss=input("Save output to file? (y/n) ")
376                if (ss == 'q'):
377                  return ss
378                if (ss != 'y'):
379                  isave=0
380                  print("Save off.")
381                else:
382                  isave=1
383                  print("Save on.")
384              
385              
386              def catchc(prompt, saveval):
387              
388                valstr=input(prompt)
389                if (valstr == 'q'):
390 tplarson 1.1     return valstr
391                while (valstr == 'c'):
392                  c=queryplotparms()
393                  if (c == 'q'):
394                    return c
395                  valstr=input(prompt)
396                if (valstr == ''):
397                  print("Using saved value %i." % saveval)
398                  valstr=str(saveval)
399                return valstr
400              
401              def printcmaps():
402              
403                print("Options include the following nonexhaustive list.  You may append '_r' to any name to reverse the map. \
404              More information is available at https://matplotlib.org/tutorials/colors/colormaps.html")
405              
406                cmaps = [('Perceptually Uniform Sequential', [
407                          'viridis', 'plasma', 'inferno', 'magma', 'cividis']),
408                       ('Sequential', [
409                          'Greys', 'Purples', 'Blues', 'Greens', 'Oranges', 'Reds',
410                          'YlOrBr', 'YlOrRd', 'OrRd', 'PuRd', 'RdPu', 'BuPu',
411 tplarson 1.1             'GnBu', 'PuBu', 'YlGnBu', 'PuBuGn', 'BuGn', 'YlGn']),
412                       ('Sequential (2)', [
413                          'binary', 'gist_yarg', 'gist_gray', 'gray', 'bone', 'pink',
414                          'spring', 'summer', 'autumn', 'winter', 'cool', 'Wistia',
415                          'hot', 'afmhot', 'gist_heat', 'copper']),
416                       ('Diverging', [
417                          'PiYG', 'PRGn', 'BrBG', 'PuOr', 'RdGy', 'RdBu',
418                          'RdYlBu', 'RdYlGn', 'Spectral', 'coolwarm', 'bwr', 'seismic']),
419              #         ('Cyclic', ['twilight', 'twilight_shifted', 'hsv']),
420                       ('Qualitative', [
421                          'Pastel1', 'Pastel2', 'Paired', 'Accent',
422                          'Dark2', 'Set1', 'Set2', 'Set3',
423                          'tab10', 'tab20', 'tab20b', 'tab20c']),
424                       ('Miscellaneous', ['hsv', 
425                          'flag', 'prism', 'ocean', 'gist_earth', 'terrain', 'gist_stern',
426                          'gnuplot', 'gnuplot2', 'CMRmap', 'cubehelix', 'brg',
427                          'gist_rainbow', 'rainbow', 'jet', 'nipy_spectral', 'gist_ncar'])]
428              
429                for c in cmaps:
430                  print("%s:" % c[0])
431                  for m in c[1]:
432 tplarson 1.1       print(m,end=" ")
433                  print("")
434              
435              
436              def nothing():
437                return None
438              
439              nframes=64
440              dpi=300

441 tplarson 1.2 icolshift=0
442              pngdirfmt = './png_out/{:s}{:s}.l{:d}m{:d}n{:d}'

443 tplarson 1.1 animate=nothing

444 tplarson 1.2 calcimage=nothing

445 tplarson 1.1 
446              def drawfigure(var, fsize=5):
447              

448 tplarson 1.2 #  below was written for single images. same functionality is now 
449              #  achieved by setting vmin and vmax.  probably exactly the same 
450              #  but new code is clearer. 
451              #  for animations.
452              #  if (icolshift != 0) and plotvar in ['Vr','Vt','Vp']:
453              #    mn,mx=var.min(),var.max()
454              #    absarr=np.abs([mn,mx])
455              #    frac=(mx-mn)/(2*absarr.max())
456              #    table=cm.get_cmap(colormap,256/frac)
457              #    if (absarr[0]>absarr[1]):
458              #      newcolors=table(np.linspace(0,frac,256))
459              #    else:
460              #      newcolors=table(np.linspace(1-frac,1,256))
461              #    newmap=ListedColormap(newcolors, name='soshtmp')
462              #    cm.register_cmap(cmap=newmap)
463              #    usemap='soshtmp'
464              #  else:
465              #    usemap=colormap
466              

467 tplarson 1.1   fig, ax = plt.subplots(num=1, figsize=(fsize, fsize))

468 tplarson 1.2   ax.clear()
469              #  im=ax.imshow(var,cmap=usemap)
470                im=ax.imshow(var, origin='lower', cmap=colormap)
471              
472                if (icolshift == 1) and plotvar in ['Vr','Vt','Vp']:
473                  mn,mx=var.min(),var.max()
474                  maxabs=np.abs([mn,mx]).max()
475              #    im=ax.imshow(var, cmap=colormap, vmin=-maxabs, vmax=maxabs)
476                  im.set_clim(vmin=-maxabs, vmax=maxabs)
477                  print("Scaling to maxval = %f"%maxabs)
478                elif (icolshift == 2):
479                  maxval=0.0
480                  for i in range(nframes):
481                    d=calcimage(i)
482                    val=np.abs(d).max()
483                    if (val > maxval):
484                      maxval=val
485                  if plotvar in ['Vr','Vt','Vp']:
486              #      im=ax.imshow(var, cmap=colormap, vmin=-maxval, vmax=maxval)
487                    im.set_clim(vmin=-maxval, vmax=maxval)
488                  else:
489 tplarson 1.2       im.set_clim(vmin=0.0, vmax=maxval)
490                  print("Scaling to maxval = %f"%maxval)
491              #  else:
492              #    im=ax.imshow(var, cmap=colormap)

493 tplarson 1.1 
494                ax.set_axis_off()
495                fig.canvas.draw()
496              
497                return (fig,im)
498              

499 tplarson 1.2 def savefigure(ianimate=1, label=''):

500 tplarson 1.1 
501                l=lsave
502                m=msave
503                n=nsave

504 tplarson 1.2   if (ianimate == 0):
505                  savestr='png_out/'+plotvar+label+'.l%im%in%i.png' % (l,m,n)

506 tplarson 1.1     plt.savefig(savestr, dpi=dpi)
507                  print("File saved.")
508                else:

509 tplarson 1.2     pngdir = pngdirfmt.format(plotvar, label, l, m, n)

510 tplarson 1.1     if not os.path.exists(pngdir):
511                    os.makedirs(pngdir)
512                  print("Writing files...", end="")
513                  for i in range(nframes):

514 tplarson 1.2       print(i, end=" ", flush=True)

515 tplarson 1.1       animate(i)
516                    fpath = os.path.join(pngdir, '{:03d}.png'.format(i))
517                    plt.savefig(fpath, dpi=dpi)
518                  print("done.")
519              
520