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              
  7              modeln=np.array([])
  8              modell=np.array([])
  9              modelnu=np.array([])
 10              rmesh=np.array([])
 11              rho=np.array([])
 12              R=0.0
 13              
 14              def loadmodel():
 15              
 16              #  in_dir = '/home/tplarson/solar/mods'
 17                in_dir = 'mods'
 18                fname = 'mods_p3_eigfcn.h5'
 19                hf = h5py.File(os.path.join(in_dir, fname))
 20              
 21              # load r, rho and R from "model"
 22 tplarson 1.1   global rmesh, rho, R
 23                rmesh, rho, R = [hf['model'][k].value for k in ['r', 'rho', 'R']]
 24                rmesh /= R
 25              
 26              # load l and n from "modes"
 27                global modeln, modell, modelnu
 28                modeln, modell, modelnu = [hf['modes'][k].value for k in ['n', 'l', 'nu']]
 29              
 30              
 31              def getradial(l,n):
 32              
 33              #  in_dir = '/home/tplarson/solar/mods'
 34                in_dir = 'mods'
 35                fname = 'mods_p3_eigfcn.h5'
 36                hf = h5py.File(os.path.join(in_dir, fname))
 37              
 38              # open y1 and y2 data sets from "modes" (without reading the data)
 39                y1ds = hf['modes/y1'] 
 40                y2ds = hf['modes/y2']
 41              
 42                idx = ((modeln == n) & (modell == l))
 43 tplarson 1.1   y1, y2 = y1ds[idx,:], y2ds[idx,:]
 44              
 45              # compute xi_r and xi_h
 46                L2 = l * (l + 1)
 47                xir = (y1 * R).flatten()
 48                if (l > 0):
 49                  xih = (y2 * R / L2).flatten()
 50                else:
 51                  xih = 0.0*xir
 52              
 53                return (xir, xih)
 54              
 55              
 56              def writemodel(lmin=0, lmax=300, rsurf=1.0):
 57              
 58                outfile='modes.model.txt'
 59                file=open(outfile,'w')
 60                in_dir = 'mods'
 61                fname = 'mods_p3_eigfcn.h5'
 62                hf = h5py.File(os.path.join(in_dir, fname))
 63              
 64 tplarson 1.1 # load r, rho and R from "model"
 65                rmesh, c, R = [hf['model'][k].value for k in ['r', 'c', 'R']]
 66                indsurf = np.abs(rmesh-rsurf*R).argmin()
 67              
 68              # load l and n from "modes"
 69                modeln, modell, modelnu, modelsig2, modelE = [hf['modes'][k].value for k in ['n', 'l', 'nu', 'sigma2', 'E']]
 70              
 71                outfmt='{:d} {:d} {:f} {:f} {:f} {:f} {:f} {:f} {:f}'
 72                l=0
 73                while (l <= lmax):
 74                  ind = (modell == l)
 75                  nlist = modeln[ind]
 76                  for n in nlist:
 77                    xir, xih = getradial(l,n)
 78              #      rc=float(xir[indsurf])
 79              #      hc=float(xih[indsurf])
 80                    rc=np.interp(rsurf*R,rmesh,xir)
 81                    hc=np.interp(rsurf*R,rmesh,xih)
 82                    ind2 = (modell == l) & (modeln == n)
 83                    nu = float(modelnu[ind2])
 84                    erg = float(modelE[ind2])*1e6
 85 tplarson 1.1       sig2 = float(modelsig2[ind2])
 86                    if (l > 0):
 87                      L=np.sqrt(l*(l+1))
 88              #        indturn=np.abs(rmesh/c - L/(2*np.pi*nu)).argmin()
 89              #        rturn=float(rmesh[indturn]/R)
 90                      rturn=np.interp(L/(2*np.pi*nu),rmesh/c,rmesh)/R
 91                    else:
 92                      rturn=0.0
 93                    amp=nu*np.sqrt(np.square(rc)+l*(l+1)*np.square(hc))
 94                    nu*=1e6
 95                    amp=float(amp)/30e6
 96              #      outstr=f'{l:d} {n:d} {nu:f} {amp:f} {sig2:f} {rc:f} {hc:f} {rturn:f} {erg:f}'
 97                    outstr=outfmt.format(int(l),int(n),nu,amp,sig2,rc,hc,rturn,erg)
 98                    file.write('%s\n' % (outstr))
 99                  l+=1
100                file.close()
101              
102              
103              
104              def image2sphere(xpixels=1000, ypixels=1000, distobs=220.0, \
105                               bangle=-30.0, pangle=0.0):
106 tplarson 1.1 
107                p = pangle*np.pi/180
108                b0 = bangle*np.pi/180
109                #distobs = 220.0  # solar radii
110                #xpixels = 1000
111                #ypixels = 1000 
112                x0 = xpixels/2-0.5
113                y0 = ypixels/2-0.5
114                imscale = 1.97784*1024/xpixels
115                scale = imscale/(180*60*60/np.pi)
116                rsun=np.tan(np.arcsin(1/distobs))/scale
117              
118              # Sun is sitting at the center of the main coordinate system and has radius 1.
119              # Observer is at robs=(xobs,yobs,zobs) moving with a velocity vobs.
120              # Start by setting robs from distobs and b0
121              
122                robs_x = distobs * np.cos(b0)
123                robs_y = 0.0
124                robs_z = distobs * np.sin(b0)
125              
126              # image coordinates
127 tplarson 1.1   xx = np.linspace(0, xpixels-1, xpixels)
128                yy = np.linspace(0, ypixels-1, ypixels)
129                xx, yy = np.meshgrid(xx, yy)
130              
131                x2 = scale*(xx - x0)
132                y2 = scale*(yy - y0)
133              # Rotate by the P-angle. New coordinate system has the y-axis pointing
134              # towards the solar north pole.
135                x1 =  x2*np.cos(p) + y2*np.sin(p)
136                y1 = -x2*np.sin(p) + y2*np.cos(p)
137              
138              # Now transform to put the coordinates into the solar coordinate system.
139              # First find the directions (vecx and vecy) the x2 and y2 coordinate
140              # axis correspond to. vecsun points towards the Sun. Note that the
141              # (x2,y2,Sun) system is left handed. These vectors are unit vectors.
142              
143                vecx_x=0.0
144                vecx_y=1.0
145                vecx_z=0.0
146                vecy_x=-np.sin(b0)
147                vecy_y=0.0
148 tplarson 1.1   vecy_z=np.cos(b0)
149                vecsun_x=-np.cos(b0)
150                vecsun_y=0.0
151                vecsun_z=-np.sin(b0)
152              
153              # Now the proper direction can be found. These are not unit vectors.
154                x = vecx_x*x1 + vecy_x*y1 + vecsun_x
155                y = vecx_y*x1 + vecy_y*y1 + vecsun_y
156                z = vecx_z*x1 + vecy_z*y1 + vecsun_z
157                qq = 1/np.sqrt(x*x + y*y + z*z)
158              # Make them unit vectors.
159                x*=qq
160                y*=qq
161                z*=qq
162              
163              # Now find intersection with the Sun.
164              # Solve quadratic equation |robs+q*[x1,y1,z1]|=1 for q
165              # a, b and c are terms in a*x^2+bx+c=0. a==1 since [x1,y1,z1] is unit vector.
166                c = robs_x*robs_x + robs_y*robs_y + robs_z*robs_z -1
167                b = 2*(x*robs_x+y*robs_y+z*robs_z)
168                d = b*b - 4*c
169 tplarson 1.1   index = (d >= 0)
170                q = np.zeros([xpixels,ypixels])
171                xsun = np.zeros([xpixels,ypixels])
172                ysun = np.zeros([xpixels,ypixels])
173                zsun = np.zeros([xpixels,ypixels])
174                q[index]=(-b[index] - np.sqrt(d[index]))/2
175                xsun[index]=robs_x + x[index]*q[index]
176                ysun[index]=robs_y + y[index]*q[index]
177                zsun[index]=robs_z + z[index]*q[index]
178              
179                phisun = np.arctan2(ysun,xsun)
180                thetasun = np.pi/2 - np.arcsin(zsun)
181              
182                ph=np.ma.array(phisun, mask=np.logical_not(index))
183                th=np.ma.array(thetasun, mask=np.logical_not(index))
184              
185                return (ph, th) 
186              
187              
188              def image2rtheta(xpixels=1000, ypixels=1000, distobs=220.0):
189              
190 tplarson 1.1   imscale = 1.97784*1024/xpixels
191                scale = imscale/(180*60*60/np.pi)
192                rsun=np.tan(np.arcsin(1/distobs))/scale
193                x0 = xpixels/2-0.5
194                y0 = ypixels/2-0.5
195                xx = (np.linspace(0, xpixels-1, xpixels)-x0)/rsun
196                yy = (np.linspace(0, ypixels-1, ypixels)-y0)/rsun
197                xx, yy = np.meshgrid(xx, yy)
198                rr = np.sqrt(xx*xx+yy*yy)
199                index = (rr <= 1.0)
200                lat = np.arctan(yy/np.abs(xx))
201                r = np.ma.array(rr, mask=np.logical_not(index))
202                theta = np.ma.array(np.pi/2 - lat, mask=np.logical_not(index))
203              
204                return (r, theta)
205              
206              
207              def setplm(l, m, x, plm, dplm): 
208              
209              # adapted from setplm.pro by Jesper Schou
210              # Set plm(*,l)=P_l^m (x) for l=m,lmax
211 tplarson 1.1 # optionally sets dplm(*,l)={{dP_l^m} \over dx} (x) for l=m,lmax
212              # P_l^m 's are normalized to have \int_{-1}^1 (P_l^m (x))^2 dx = 1
213              # Works up to l \approx 1800, see ~/invnew/plm.f for details
214              
215                eps = 1.0e-12
216                x1 = np.maximum(np.minimum(x,1-eps),eps-1)
217                x2 = 1.0/(x1*x1-1.0)
218                x3 = x1*x2
219                c = np.sqrt((2*m+1)/2.0)
220                for i in range(1,m+1):
221                  c *= -np.sqrt(1.0-0.5/i)
222                plm[...,0] = c*(np.sqrt(1.0-x1*x1))**m
223                if (l > m):
224                  c = np.sqrt(2.0*m+3.0)
225                  plm[...,1] = c*x1*plm[...,0]
226                i = m+2
227                while (i <= l):
228                  c1 = np.sqrt((4.0*i*i-1.0)/(i*i-m*m))
229                  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)))
230                  plm[...,i-m] = c1*x1*plm[...,i-m-1] - c2*plm[...,i-m-2]
231                  i+=1
232 tplarson 1.1 
233                dplm[...,0] = m*x3*plm[...,0]
234                i = m+1
235                while (i <= l):
236                  c = np.sqrt((2.0*i+1.0)*(i*i-m*m)/(2.0*i-1))
237                  dplm[...,i-m] = i*x3*plm[...,i-m] - c*x2*plm[...,i-m-1]
238                  i+=1
239              
240                return (plm[...,l-m],dplm[...,l-m])
241              
242              
243              def apols(lin, amaxin):
244              
245              # adapted from IDL procedure apols.pro written by Jesper Schou
246              
247                l=np.long(lin)
248                amax=np.minimum(amaxin,2*l)
249              
250                pols=np.zeros((2*l+1,amaxin+1))
251              # It is ESSENTIAL that x is set up such that x(-m)=-x(m) to full machine
252              # accuracy or that the re-orthogonalization is done with respect to all
253 tplarson 1.1 # previous polynomials (second option below).
254              # x=(dindgen(2*l+1)-l)/l
255                x=np.linspace(-1,1,2*l+1)
256              
257                pols[:,0]=1/np.sqrt(2*l+1.0)
258                if (amax >= 1):
259                  pols[:,1]=x/np.sqrt((l+1.0)*(2*l+1.0)/(3.0*l))
260              # for n=2l,amax do begin
261              # Set up polynomials using exact recurrence relation.
262                for n in range(2,amax+1):
263                  a0=2.0*l*(2*n-1.0)/(n*(2*l-n+1))
264                  b0=-(n-1.0)/n*(2*l+n)/(2*l-n+1)
265                  a=a0*np.sqrt((2*n+1.0)/(2*n-1.0)*(2*l-n+1.0)/(2*l+n+1.0))
266                  b=b0*np.sqrt((2*n+1.0)/(2*n-3.0)*(2*l-n+1.0)*(2*l-n+2.0)/(2*l+n+1.0)/(2*l+n))
267                  help=a*x*pols[:,n-1]+b*pols[:,n-2]
268              # Turns out that roundoff kills the algorithm, so we have to re-orthogonalize.
269              # Two choices here. First one is twice as fast and more accurate, but
270              # depends on the rounding being done in the standard IEEE way.
271              #  for j=n-2,0,-2 do begin 
272                  for j in range(n-2,-1,-2):    
273              # This choice is robust to the roundoff, but twice as slow and generally
274 tplarson 1.1 # somewhat less accurate
275              # for j=n-1,0,-1 do begin
276                    help=help-pols[:,j]*np.sum(help*pols[:,j])
277              #  end
278              # Reset norm to 1.
279                  pols[:,n]=help/np.sqrt(np.sum(np.square(help)))
280              
281              # Reset polynomials to have P_l(l)=l by setting overall norm.
282              # Note that this results in more accurate overall normalization, but
283              # that P_l(l) may be very far from l. This is the key to the algorithm.
284                c=l**2*(2*l+1.0)
285              # for n=0l,amax do begin
286                for n in range(0,amax+1):
287                  c=c*(2*l+n+1.0)/(2*l-n+1.0)
288                  pols[:,n]=pols[:,n]*np.sqrt(c/(2*n+1))
289              
290                return pols
291              
292              
293              lsave=0
294              msave=0
295 tplarson 1.1 nsave=1
296              
297              def querylmn():
298              
299                global lsave, msave, nsave
300                lstr=catchc("Enter spherical harmonic degree (l): ",lsave)
301                if lstr == 'q':
302                  return (-1,-1,-1)
303                else:
304                  while True:
305                    try:
306                      lval=int(lstr)
307                      if lval < 0:
308                        print("Degree must be greater than or equal to zero. Try again.")
309                        lstr=catchc("Enter spherical harmonic degree (l): ",lsave)
310                        if lstr == 'q':
311                          return (-1,-1,-1)
312                        continue
313                      break
314                    except:
315                      print("Invalid number, try again.")
316 tplarson 1.1         lstr=catchc("Enter spherical harmonic degree (l): ", lsave)
317                      if lstr == 'q':
318                        return (-1,-1,-1)
319              
320                mstr=catchc("Enter azimuthal order (m): ",msave)
321                if mstr == 'q':
322                  return (-1,-1,-1)
323                else:
324                  while True:
325                    try:
326                      mval=int(mstr)
327                      if np.abs(mval) > lval:
328                        print("Azimuthal order must be in the range -l <= m <= l. Try again.")
329                        mstr=catchc("Enter azimuthal order (m): ",msave)
330                        if mstr == 'q':
331                          return (-1,-1,-1)
332                        continue
333                      break
334                    except:
335                      print("Invalid number, try again.")
336                      mstr=catchc("Enter azimuthal order (m): ",msave)
337 tplarson 1.1         if mstr == 'q':
338                        return (-1,-1,-1)
339              
340                nstr=catchc("Enter radial order (n): ",nsave)
341                if nstr == 'q':
342                  return (-1,-1,-1)
343                else:
344                  while True:
345                    try:
346                      nval=int(nstr)
347                      if nval < 0:
348                        print("Radial order must be greater than or equal to zero. Try again.")
349                        nstr=catchc("Enter radial order (n): ",nsave)
350                        if nstr == 'q':
351                          return (-1,-1,-1)
352                        continue
353                      break
354                    except:
355                      print("Invalid number, try again.")
356                      nstr=catchc("Enter radial order (n): ",nsave)
357                      if nstr == 'q':
358 tplarson 1.1           return (-1,-1,-1)
359              
360                lsave=lval
361                msave=mval
362                nsave=nval
363                return (lval,mval,nval)
364              
365              
366              varlist=['Vr', 'Vt', 'Vp', 'Vh', 'Vmag', 'Vsq']
367              colormap="seismic"
368              plotvar="Vr"
369              isave=0
370              
371              def queryplotparms():
372              
373                global colormap, plotvar, isave
374                print("You may enter 'l' to list options.")
375                cmap=input("Enter name of colormap: ")
376                while True:
377                  if (cmap == 'q'):
378                    return cmap
379 tplarson 1.1     if (cmap == ''):
380                    print("Using saved value colormap = %s." % colormap)
381                    cmap=colormap
382                    break
383                  if (cmap == 'l'):
384                    printcmaps()
385                    print("Current colormap is %s." % colormap)
386                    cmap=input("Enter name of colormap: ")
387                    continue
388                  elif (cmap not in plt.colormaps()):
389                    print("That colormap not registered, try again.")
390                    cmap=input("Enter name of colormap: ")
391                    continue
392                  else:
393                    break
394                colormap=cmap
395                  
396                pvar=input("Enter variable to plot: ")
397                while True:
398                  if (pvar == 'q'):
399                    return pvar
400 tplarson 1.1     if (pvar == ''):
401                    print("Using saved value plotvar = %s." % plotvar)
402                    pvar=plotvar
403                    break
404                  if (pvar == 'l'):
405                    print("Available plotting variables are: ", end="")
406                    for i in varlist[:len(varlist)-1]:
407                      print(i, end=", ")
408                    print("and %s" % varlist[len(varlist)-1], end=".\n")
409                    print("Currently plotting %s." % plotvar)
410                    pvar=input("Enter variable to plot: ")
411                    continue
412                  elif (pvar not in varlist):
413                    print("That variable not available, try again.")
414                    pvar=input("Enter variable to plot: ")
415                    continue
416                  else:
417                    break
418                plotvar=pvar
419              
420                ss=input("Save output to file? (y/n) ")
421 tplarson 1.1   if (ss == 'q'):
422                  return ss
423                if (ss != 'y'):
424                  isave=0
425                  print("Save off.")
426                else:
427                  isave=1
428                  print("Save on.")
429              
430              
431              def catchc(prompt, saveval):
432              
433                valstr=input(prompt)
434                if (valstr == 'q'):
435                  return valstr
436                while (valstr == 'c'):
437                  c=queryplotparms()
438                  if (c == 'q'):
439                    return c
440                  valstr=input(prompt)
441                if (valstr == ''):
442 tplarson 1.1     print("Using saved value %i." % saveval)
443                  valstr=str(saveval)
444                return valstr
445              
446              def printcmaps():
447              
448                print("Options include the following nonexhaustive list.  You may append '_r' to any name to reverse the map. \
449              More information is available at https://matplotlib.org/tutorials/colors/colormaps.html")
450              
451                cmaps = [('Perceptually Uniform Sequential', [
452                          'viridis', 'plasma', 'inferno', 'magma', 'cividis']),
453                       ('Sequential', [
454                          'Greys', 'Purples', 'Blues', 'Greens', 'Oranges', 'Reds',
455                          'YlOrBr', 'YlOrRd', 'OrRd', 'PuRd', 'RdPu', 'BuPu',
456                          'GnBu', 'PuBu', 'YlGnBu', 'PuBuGn', 'BuGn', 'YlGn']),
457                       ('Sequential (2)', [
458                          'binary', 'gist_yarg', 'gist_gray', 'gray', 'bone', 'pink',
459                          'spring', 'summer', 'autumn', 'winter', 'cool', 'Wistia',
460                          'hot', 'afmhot', 'gist_heat', 'copper']),
461                       ('Diverging', [
462                          'PiYG', 'PRGn', 'BrBG', 'PuOr', 'RdGy', 'RdBu',
463 tplarson 1.1             'RdYlBu', 'RdYlGn', 'Spectral', 'coolwarm', 'bwr', 'seismic']),
464              #         ('Cyclic', ['twilight', 'twilight_shifted', 'hsv']),
465                       ('Qualitative', [
466                          'Pastel1', 'Pastel2', 'Paired', 'Accent',
467                          'Dark2', 'Set1', 'Set2', 'Set3',
468                          'tab10', 'tab20', 'tab20b', 'tab20c']),
469                       ('Miscellaneous', ['hsv', 
470                          'flag', 'prism', 'ocean', 'gist_earth', 'terrain', 'gist_stern',
471                          'gnuplot', 'gnuplot2', 'CMRmap', 'cubehelix', 'brg',
472                          'gist_rainbow', 'rainbow', 'jet', 'nipy_spectral', 'gist_ncar'])]
473              
474                for c in cmaps:
475                  print("%s:" % c[0])
476                  for m in c[1]:
477                    print(m,end=" ")
478                  print("")
479              
480              
481              def nothing():
482                return None
483              
484 tplarson 1.1 nframes=64
485              dpi=300
486              pngdirfmt = './png/l{:d}_m{:d}_n{:d}_{:s}'
487              animate=nothing
488              
489              def drawfigure(var, fsize=5):
490              
491                fig, ax = plt.subplots(num=1, figsize=(fsize, fsize))
492              #  ax.set(xlim=(-lim, lim), ylim=(-lim, lim))
493              
494                ax.clear()
495                im=ax.imshow(var,cmap=colormap)
496                ax.set_axis_off()
497                fig.canvas.draw()
498              
499                return (fig,im)
500              
501              def savefigure(animateflag=1):
502              
503                l=lsave
504                m=msave
505 tplarson 1.1   n=nsave
506                if (animateflag == 0):
507                  savestr="l%im%in%i.png" % (l,m,n)
508                  plt.savefig(savestr, dpi=dpi)
509                  print("File saved.")
510                else:
511                  pngdir = pngdirfmt.format(l, m, n, plotvar)
512                  if not os.path.exists(pngdir):
513                    os.makedirs(pngdir)
514                  print("Writing files...", end="")
515                  for i in range(nframes):
516                    animate(i)
517                    fpath = os.path.join(pngdir, '{:03d}.png'.format(i))
518                    print(i, end=" ", flush=True)
519                    plt.savefig(fpath, dpi=dpi)
520                  print("done.")
521              
522