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 tplarson 1.3 lsave=[0]
249 msave=[0]
250 nsave=[1]
|
251 tplarson 1.1
|
252 tplarson 1.3 def querylmn(index):
|
253 tplarson 1.1
254 global lsave, msave, nsave
|
255 tplarson 1.3 if (index < 0):
256 if (index < -len(lsave)):
257 i=-len(lsave)
258 else:
259 i=index
260 else:
261 i=-1
262 lstr=catchc("Enter spherical harmonic degree (l): ",lsave[i])
|
263 tplarson 1.1 if lstr == 'q':
264 return (-1,-1,-1)
265 else:
266 while True:
267 try:
268 lval=int(lstr)
269 if lval < 0:
270 print("Degree must be greater than or equal to zero. Try again.")
|
271 tplarson 1.3 lstr=catchc("Enter spherical harmonic degree (l): ",lsave[i])
|
272 tplarson 1.1 if lstr == 'q':
273 return (-1,-1,-1)
274 continue
275 break
276 except:
277 print("Invalid number, try again.")
|
278 tplarson 1.3 lstr=catchc("Enter spherical harmonic degree (l): ", lsave[i])
|
279 tplarson 1.1 if lstr == 'q':
280 return (-1,-1,-1)
281
|
282 tplarson 1.3 mstr=catchc("Enter azimuthal order (m): ",msave[i])
|
283 tplarson 1.1 if mstr == 'q':
284 return (-1,-1,-1)
285 else:
286 while True:
287 try:
288 mval=int(mstr)
289 if np.abs(mval) > lval:
290 print("Azimuthal order must be in the range -l <= m <= l. Try again.")
|
291 tplarson 1.3 mstr=catchc("Enter azimuthal order (m): ",msave[i])
|
292 tplarson 1.1 if mstr == 'q':
293 return (-1,-1,-1)
294 continue
295 break
296 except:
297 print("Invalid number, try again.")
|
298 tplarson 1.3 mstr=catchc("Enter azimuthal order (m): ",msave[i])
|
299 tplarson 1.1 if mstr == 'q':
300 return (-1,-1,-1)
301
|
302 tplarson 1.3 nstr=catchc("Enter radial order (n): ",nsave[i])
|
303 tplarson 1.1 if nstr == 'q':
304 return (-1,-1,-1)
305 else:
306 while True:
307 try:
308 nval=int(nstr)
309 if nval < 0:
310 print("Radial order must be greater than or equal to zero. Try again.")
|
311 tplarson 1.3 nstr=catchc("Enter radial order (n): ",nsave[i])
|
312 tplarson 1.1 if nstr == 'q':
313 return (-1,-1,-1)
314 continue
315 break
316 except:
317 print("Invalid number, try again.")
|
318 tplarson 1.3 nstr=catchc("Enter radial order (n): ",nsave[i])
|
319 tplarson 1.1 if nstr == 'q':
320 return (-1,-1,-1)
321
|
322 tplarson 1.3 lsave.append(lval)
323 msave.append(mval)
324 nsave.append(nval)
|
325 tplarson 1.1 return (lval,mval,nval)
326
327
328 varlist=['Vr', 'Vt', 'Vp', 'Vh', 'Vmag', 'Vsq']
329 colormap="seismic"
330 plotvar="Vr"
331 isave=0
332
333 def queryplotparms():
334
335 global colormap, plotvar, isave
336 print("You may enter 'l' to list options.")
337 cmap=input("Enter name of colormap: ")
338 while True:
339 if (cmap == 'q'):
340 return cmap
341 if (cmap == ''):
342 print("Using saved value colormap = %s." % colormap)
343 cmap=colormap
344 break
345 if (cmap == 'l'):
346 tplarson 1.1 printcmaps()
347 print("Current colormap is %s." % colormap)
348 cmap=input("Enter name of colormap: ")
349 continue
350 elif (cmap not in plt.colormaps()):
351 print("That colormap not registered, try again.")
352 cmap=input("Enter name of colormap: ")
353 continue
354 else:
355 break
356 colormap=cmap
357
358 pvar=input("Enter variable to plot: ")
359 while True:
360 if (pvar == 'q'):
361 return pvar
362 if (pvar == ''):
363 print("Using saved value plotvar = %s." % plotvar)
364 pvar=plotvar
365 break
366 if (pvar == 'l'):
367 tplarson 1.1 print("Available plotting variables are: ", end="")
368 for i in varlist[:len(varlist)-1]:
369 print(i, end=", ")
370 print("and %s" % varlist[len(varlist)-1], end=".\n")
371 print("Currently plotting %s." % plotvar)
372 pvar=input("Enter variable to plot: ")
373 continue
374 elif (pvar not in varlist):
375 print("That variable not available, try again.")
376 pvar=input("Enter variable to plot: ")
377 continue
378 else:
379 break
380 plotvar=pvar
381
382 ss=input("Save output to file? (y/n) ")
383 if (ss == 'q'):
384 return ss
385 if (ss != 'y'):
386 isave=0
387 print("Save off.")
388 tplarson 1.1 else:
389 isave=1
390 print("Save on.")
391
392
393 def catchc(prompt, saveval):
394
395 valstr=input(prompt)
396 if (valstr == 'q'):
397 return valstr
398 while (valstr == 'c'):
399 c=queryplotparms()
400 if (c == 'q'):
401 return c
402 valstr=input(prompt)
403 if (valstr == ''):
|
404 tplarson 1.3 print("Using saved value %s." % str(saveval))
|
405 tplarson 1.1 valstr=str(saveval)
406 return valstr
407
408 def printcmaps():
409
410 print("Options include the following nonexhaustive list. You may append '_r' to any name to reverse the map. \
411 More information is available at https://matplotlib.org/tutorials/colors/colormaps.html")
412
413 cmaps = [('Perceptually Uniform Sequential', [
414 'viridis', 'plasma', 'inferno', 'magma', 'cividis']),
415 ('Sequential', [
416 'Greys', 'Purples', 'Blues', 'Greens', 'Oranges', 'Reds',
417 'YlOrBr', 'YlOrRd', 'OrRd', 'PuRd', 'RdPu', 'BuPu',
418 'GnBu', 'PuBu', 'YlGnBu', 'PuBuGn', 'BuGn', 'YlGn']),
419 ('Sequential (2)', [
420 'binary', 'gist_yarg', 'gist_gray', 'gray', 'bone', 'pink',
421 'spring', 'summer', 'autumn', 'winter', 'cool', 'Wistia',
422 'hot', 'afmhot', 'gist_heat', 'copper']),
423 ('Diverging', [
424 'PiYG', 'PRGn', 'BrBG', 'PuOr', 'RdGy', 'RdBu',
425 'RdYlBu', 'RdYlGn', 'Spectral', 'coolwarm', 'bwr', 'seismic']),
426 tplarson 1.1 # ('Cyclic', ['twilight', 'twilight_shifted', 'hsv']),
427 ('Qualitative', [
428 'Pastel1', 'Pastel2', 'Paired', 'Accent',
429 'Dark2', 'Set1', 'Set2', 'Set3',
430 'tab10', 'tab20', 'tab20b', 'tab20c']),
431 ('Miscellaneous', ['hsv',
432 'flag', 'prism', 'ocean', 'gist_earth', 'terrain', 'gist_stern',
433 'gnuplot', 'gnuplot2', 'CMRmap', 'cubehelix', 'brg',
434 'gist_rainbow', 'rainbow', 'jet', 'nipy_spectral', 'gist_ncar'])]
435
436 for c in cmaps:
437 print("%s:" % c[0])
438 for m in c[1]:
439 print(m,end=" ")
440 print("")
441
442
443 def nothing():
444 return None
445
446 nframes=64
447 tplarson 1.1 dpi=300
|
448 tplarson 1.2 icolshift=0
449 pngdirfmt = './png_out/{:s}{:s}.l{:d}m{:d}n{:d}'
|
450 tplarson 1.1 animate=nothing
|
451 tplarson 1.2 calcimage=nothing
|
452 tplarson 1.1
453 def drawfigure(var, fsize=5):
454
|
455 tplarson 1.2 # below was written for single images. same functionality is now
456 # achieved by setting vmin and vmax. probably exactly the same
457 # but new code is clearer.
458 # for animations.
459 # if (icolshift != 0) and plotvar in ['Vr','Vt','Vp']:
460 # mn,mx=var.min(),var.max()
461 # absarr=np.abs([mn,mx])
462 # frac=(mx-mn)/(2*absarr.max())
463 # table=cm.get_cmap(colormap,256/frac)
464 # if (absarr[0]>absarr[1]):
465 # newcolors=table(np.linspace(0,frac,256))
466 # else:
467 # newcolors=table(np.linspace(1-frac,1,256))
468 # newmap=ListedColormap(newcolors, name='soshtmp')
469 # cm.register_cmap(cmap=newmap)
470 # usemap='soshtmp'
471 # else:
472 # usemap=colormap
473
|
474 tplarson 1.1 fig, ax = plt.subplots(num=1, figsize=(fsize, fsize))
|
475 tplarson 1.2 ax.clear()
476 # im=ax.imshow(var,cmap=usemap)
477 im=ax.imshow(var, origin='lower', cmap=colormap)
478
479 if (icolshift == 1) and plotvar in ['Vr','Vt','Vp']:
480 mn,mx=var.min(),var.max()
481 maxabs=np.abs([mn,mx]).max()
482 # im=ax.imshow(var, cmap=colormap, vmin=-maxabs, vmax=maxabs)
483 im.set_clim(vmin=-maxabs, vmax=maxabs)
484 print("Scaling to maxval = %f"%maxabs)
485 elif (icolshift == 2):
486 maxval=0.0
487 for i in range(nframes):
488 d=calcimage(i)
489 val=np.abs(d).max()
490 if (val > maxval):
491 maxval=val
492 if plotvar in ['Vr','Vt','Vp']:
493 # im=ax.imshow(var, cmap=colormap, vmin=-maxval, vmax=maxval)
494 im.set_clim(vmin=-maxval, vmax=maxval)
495 else:
496 tplarson 1.2 im.set_clim(vmin=0.0, vmax=maxval)
497 print("Scaling to maxval = %f"%maxval)
498 # else:
499 # im=ax.imshow(var, cmap=colormap)
|
500 tplarson 1.1
501 ax.set_axis_off()
502 fig.canvas.draw()
503
504 return (fig,im)
505
|
506 tplarson 1.2 def savefigure(ianimate=1, label=''):
|
507 tplarson 1.1
|
508 tplarson 1.3 l=lsave[-1]
509 m=msave[-1]
510 n=nsave[-1]
|
511 tplarson 1.2 if (ianimate == 0):
512 savestr='png_out/'+plotvar+label+'.l%im%in%i.png' % (l,m,n)
|
513 tplarson 1.1 plt.savefig(savestr, dpi=dpi)
514 print("File saved.")
515 else:
|
516 tplarson 1.2 pngdir = pngdirfmt.format(plotvar, label, l, m, n)
|
517 tplarson 1.1 if not os.path.exists(pngdir):
518 os.makedirs(pngdir)
519 print("Writing files...", end="")
520 for i in range(nframes):
|
521 tplarson 1.2 print(i, end=" ", flush=True)
|
522 tplarson 1.1 animate(i)
523 fpath = os.path.join(pngdir, '{:03d}.png'.format(i))
524 plt.savefig(fpath, dpi=dpi)
525 print("done.")
526
527
|
Return to
sosh.py
CVS log
Up to [Development] / JSOC / proj / globalhs / sosh