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 tplarson 1.4 bangle=30.0, pangle=0.0, xoffset=0.0, yoffset=0.0):
|
109 tplarson 1.1
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 tplarson 1.4 x0 = xpixels/2-0.5 + xoffset
116 y0 = ypixels/2-0.5 + yoffset
|
117 tplarson 1.1 imscale = 1.97784*1024/xpixels
118 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 tplarson 1.1 x1 = x2*np.cos(p) + y2*np.sin(p)
139 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 tplarson 1.1 z = vecx_z*x1 + vecy_z*y1 + vecsun_z
160 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 tplarson 1.1 zsun[index]=robs_z + z[index]*q[index]
181
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 tplarson 1.4 def image2rtheta(xpixels=1000, ypixels=1000, distobs=220.0, xoffset=0.0, yoffset=0.0, \
192 bangle=0.0, pangle=0.0, gamma=0.0):
|
193 tplarson 1.1
|
194 tplarson 1.4 p = pangle*np.pi/180
195 b0 = bangle*np.pi/180
196 g = gamma*np.pi/180
|
197 tplarson 1.1 imscale = 1.97784*1024/xpixels
198 scale = imscale/(180*60*60/np.pi)
199 rsun=np.tan(np.arcsin(1/distobs))/scale
|
200 tplarson 1.4 x0 = xpixels/2-0.5 + xoffset
201 y0 = ypixels/2-0.5 + yoffset
|
202 tplarson 1.1 xx = (np.linspace(0, xpixels-1, xpixels)-x0)/rsun
|
203 tplarson 1.4 zz = (np.linspace(0, ypixels-1, ypixels)-y0)/rsun
204 xx, zz = np.meshgrid(xx, zz)
205
206 xx /= np.cos(g)
207 zz /= np.cos(b0)
208 x1 = xx*np.cos(p) + zz*np.sin(p)
209 z1 = -xx*np.sin(p) + zz*np.cos(p)
210
211 rr = np.sqrt(x1*x1 + z1*z1)
|
212 tplarson 1.1 index = (rr <= 1.0)
213 r = np.ma.array(rr, mask=np.logical_not(index))
|
214 tplarson 1.4 angle = np.arctan2(x1,z1)
215 theta = np.ma.array(np.abs(angle), mask=r.mask)
216 phi = 0.0*theta
217 index = (angle < 0.0)
218 phi[index]=np.pi
|
219 tplarson 1.1
|
220 tplarson 1.4 return (r, theta, phi)
|
221 tplarson 1.1
222
223 def setplm(l, m, x, plm, dplm):
224
225 # adapted from setplm.pro by Jesper Schou
226 # Set plm(*,l)=P_l^m (x) for l=m,lmax
227 # optionally sets dplm(*,l)={{dP_l^m} \over dx} (x) for l=m,lmax
228 # P_l^m 's are normalized to have \int_{-1}^1 (P_l^m (x))^2 dx = 1
229 # Works up to l \approx 1800, see ~/invnew/plm.f for details
230
231 eps = 1.0e-12
232 x1 = np.maximum(np.minimum(x,1-eps),eps-1)
|
233 tplarson 1.2 # x1 = x
|
234 tplarson 1.1 x2 = 1.0/(x1*x1-1.0)
235 x3 = x1*x2
236 c = np.sqrt((2*m+1)/2.0)
237 for i in range(1,m+1):
238 c *= -np.sqrt(1.0-0.5/i)
239 plm[...,0] = c*(np.sqrt(1.0-x1*x1))**m
240 if (l > m):
241 c = np.sqrt(2.0*m+3.0)
242 plm[...,1] = c*x1*plm[...,0]
243 i = m+2
244 while (i <= l):
245 c1 = np.sqrt((4.0*i*i-1.0)/(i*i-m*m))
246 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)))
247 plm[...,i-m] = c1*x1*plm[...,i-m-1] - c2*plm[...,i-m-2]
248 i+=1
249
250 dplm[...,0] = m*x3*plm[...,0]
251 i = m+1
252 while (i <= l):
253 c = np.sqrt((2.0*i+1.0)*(i*i-m*m)/(2.0*i-1))
254 dplm[...,i-m] = i*x3*plm[...,i-m] - c*x2*plm[...,i-m-1]
255 tplarson 1.1 i+=1
256
257 return (plm[...,l-m],dplm[...,l-m])
258
259
|
260 tplarson 1.3 lsave=[0]
261 msave=[0]
262 nsave=[1]
|
263 tplarson 1.1
|
264 tplarson 1.3 def querylmn(index):
|
265 tplarson 1.1
266 global lsave, msave, nsave
|
267 tplarson 1.3 if (index < 0):
268 if (index < -len(lsave)):
269 i=-len(lsave)
270 else:
271 i=index
272 else:
273 i=-1
274 lstr=catchc("Enter spherical harmonic degree (l): ",lsave[i])
|
275 tplarson 1.1 if lstr == 'q':
276 return (-1,-1,-1)
277 else:
278 while True:
279 try:
280 lval=int(lstr)
281 if lval < 0:
282 print("Degree must be greater than or equal to zero. Try again.")
|
283 tplarson 1.3 lstr=catchc("Enter spherical harmonic degree (l): ",lsave[i])
|
284 tplarson 1.1 if lstr == 'q':
285 return (-1,-1,-1)
286 continue
287 break
288 except:
289 print("Invalid number, try again.")
|
290 tplarson 1.3 lstr=catchc("Enter spherical harmonic degree (l): ", lsave[i])
|
291 tplarson 1.1 if lstr == 'q':
292 return (-1,-1,-1)
293
|
294 tplarson 1.3 mstr=catchc("Enter azimuthal order (m): ",msave[i])
|
295 tplarson 1.1 if mstr == 'q':
296 return (-1,-1,-1)
297 else:
298 while True:
299 try:
300 mval=int(mstr)
301 if np.abs(mval) > lval:
302 print("Azimuthal order must be in the range -l <= m <= l. Try again.")
|
303 tplarson 1.3 mstr=catchc("Enter azimuthal order (m): ",msave[i])
|
304 tplarson 1.1 if mstr == 'q':
305 return (-1,-1,-1)
306 continue
307 break
308 except:
309 print("Invalid number, try again.")
|
310 tplarson 1.3 mstr=catchc("Enter azimuthal order (m): ",msave[i])
|
311 tplarson 1.1 if mstr == 'q':
312 return (-1,-1,-1)
313
|
314 tplarson 1.3 nstr=catchc("Enter radial order (n): ",nsave[i])
|
315 tplarson 1.1 if nstr == 'q':
316 return (-1,-1,-1)
317 else:
318 while True:
319 try:
320 nval=int(nstr)
321 if nval < 0:
322 print("Radial order must be greater than or equal to zero. Try again.")
|
323 tplarson 1.3 nstr=catchc("Enter radial order (n): ",nsave[i])
|
324 tplarson 1.1 if nstr == 'q':
325 return (-1,-1,-1)
326 continue
327 break
328 except:
329 print("Invalid number, try again.")
|
330 tplarson 1.3 nstr=catchc("Enter radial order (n): ",nsave[i])
|
331 tplarson 1.1 if nstr == 'q':
332 return (-1,-1,-1)
333
|
334 tplarson 1.3 lsave.append(lval)
335 msave.append(mval)
336 nsave.append(nval)
|
337 tplarson 1.1 return (lval,mval,nval)
338
339
340 varlist=['Vr', 'Vt', 'Vp', 'Vh', 'Vmag', 'Vsq']
341 colormap="seismic"
342 plotvar="Vr"
343 isave=0
344
345 def queryplotparms():
346
347 global colormap, plotvar, isave
348 print("You may enter 'l' to list options.")
349 cmap=input("Enter name of colormap: ")
350 while True:
351 if (cmap == 'q'):
352 return cmap
353 if (cmap == ''):
354 print("Using saved value colormap = %s." % colormap)
355 cmap=colormap
356 break
357 if (cmap == 'l'):
358 tplarson 1.1 printcmaps()
359 print("Current colormap is %s." % colormap)
360 cmap=input("Enter name of colormap: ")
361 continue
362 elif (cmap not in plt.colormaps()):
363 print("That colormap not registered, try again.")
364 cmap=input("Enter name of colormap: ")
365 continue
366 else:
367 break
368 colormap=cmap
369
370 pvar=input("Enter variable to plot: ")
371 while True:
372 if (pvar == 'q'):
373 return pvar
374 if (pvar == ''):
375 print("Using saved value plotvar = %s." % plotvar)
376 pvar=plotvar
377 break
378 if (pvar == 'l'):
379 tplarson 1.1 print("Available plotting variables are: ", end="")
380 for i in varlist[:len(varlist)-1]:
381 print(i, end=", ")
382 print("and %s" % varlist[len(varlist)-1], end=".\n")
383 print("Currently plotting %s." % plotvar)
384 pvar=input("Enter variable to plot: ")
385 continue
386 elif (pvar not in varlist):
387 print("That variable not available, try again.")
388 pvar=input("Enter variable to plot: ")
389 continue
390 else:
391 break
392 plotvar=pvar
393
394 ss=input("Save output to file? (y/n) ")
395 if (ss == 'q'):
396 return ss
397 if (ss != 'y'):
398 isave=0
399 print("Save off.")
400 tplarson 1.1 else:
401 isave=1
402 print("Save on.")
403
404
405 def catchc(prompt, saveval):
406
407 valstr=input(prompt)
408 if (valstr == 'q'):
409 return valstr
410 while (valstr == 'c'):
411 c=queryplotparms()
412 if (c == 'q'):
413 return c
414 valstr=input(prompt)
415 if (valstr == ''):
|
416 tplarson 1.3 print("Using saved value %s." % str(saveval))
|
417 tplarson 1.1 valstr=str(saveval)
418 return valstr
419
420 def printcmaps():
421
422 print("Options include the following nonexhaustive list. You may append '_r' to any name to reverse the map. \
423 More information is available at https://matplotlib.org/tutorials/colors/colormaps.html")
424
425 cmaps = [('Perceptually Uniform Sequential', [
426 'viridis', 'plasma', 'inferno', 'magma', 'cividis']),
427 ('Sequential', [
428 'Greys', 'Purples', 'Blues', 'Greens', 'Oranges', 'Reds',
429 'YlOrBr', 'YlOrRd', 'OrRd', 'PuRd', 'RdPu', 'BuPu',
430 'GnBu', 'PuBu', 'YlGnBu', 'PuBuGn', 'BuGn', 'YlGn']),
431 ('Sequential (2)', [
432 'binary', 'gist_yarg', 'gist_gray', 'gray', 'bone', 'pink',
433 'spring', 'summer', 'autumn', 'winter', 'cool', 'Wistia',
434 'hot', 'afmhot', 'gist_heat', 'copper']),
435 ('Diverging', [
436 'PiYG', 'PRGn', 'BrBG', 'PuOr', 'RdGy', 'RdBu',
437 'RdYlBu', 'RdYlGn', 'Spectral', 'coolwarm', 'bwr', 'seismic']),
438 tplarson 1.1 # ('Cyclic', ['twilight', 'twilight_shifted', 'hsv']),
439 ('Qualitative', [
440 'Pastel1', 'Pastel2', 'Paired', 'Accent',
441 'Dark2', 'Set1', 'Set2', 'Set3',
442 'tab10', 'tab20', 'tab20b', 'tab20c']),
443 ('Miscellaneous', ['hsv',
444 'flag', 'prism', 'ocean', 'gist_earth', 'terrain', 'gist_stern',
445 'gnuplot', 'gnuplot2', 'CMRmap', 'cubehelix', 'brg',
446 'gist_rainbow', 'rainbow', 'jet', 'nipy_spectral', 'gist_ncar'])]
447
448 for c in cmaps:
449 print("%s:" % c[0])
450 for m in c[1]:
451 print(m,end=" ")
452 print("")
453
454
455 def nothing():
456 return None
457
458 nframes=64
459 tplarson 1.1 dpi=300
|
460 tplarson 1.2 icolshift=0
461 pngdirfmt = './png_out/{:s}{:s}.l{:d}m{:d}n{:d}'
|
462 tplarson 1.1 animate=nothing
|
463 tplarson 1.2 calcimage=nothing
|
464 tplarson 1.1
465 def drawfigure(var, fsize=5):
466
|
467 tplarson 1.2 # below was written for single images. same functionality is now
468 # achieved by setting vmin and vmax. probably exactly the same
469 # but new code is clearer.
470 # for animations.
471 # if (icolshift != 0) and plotvar in ['Vr','Vt','Vp']:
472 # mn,mx=var.min(),var.max()
473 # absarr=np.abs([mn,mx])
474 # frac=(mx-mn)/(2*absarr.max())
475 # table=cm.get_cmap(colormap,256/frac)
476 # if (absarr[0]>absarr[1]):
477 # newcolors=table(np.linspace(0,frac,256))
478 # else:
479 # newcolors=table(np.linspace(1-frac,1,256))
480 # newmap=ListedColormap(newcolors, name='soshtmp')
481 # cm.register_cmap(cmap=newmap)
482 # usemap='soshtmp'
483 # else:
484 # usemap=colormap
485
|
486 tplarson 1.1 fig, ax = plt.subplots(num=1, figsize=(fsize, fsize))
|
487 tplarson 1.2 ax.clear()
488 # im=ax.imshow(var,cmap=usemap)
489 im=ax.imshow(var, origin='lower', cmap=colormap)
490
491 if (icolshift == 1) and plotvar in ['Vr','Vt','Vp']:
492 mn,mx=var.min(),var.max()
493 maxabs=np.abs([mn,mx]).max()
494 # im=ax.imshow(var, cmap=colormap, vmin=-maxabs, vmax=maxabs)
495 im.set_clim(vmin=-maxabs, vmax=maxabs)
496 print("Scaling to maxval = %f"%maxabs)
497 elif (icolshift == 2):
498 maxval=0.0
499 for i in range(nframes):
500 d=calcimage(i)
501 val=np.abs(d).max()
502 if (val > maxval):
503 maxval=val
504 if plotvar in ['Vr','Vt','Vp']:
505 # im=ax.imshow(var, cmap=colormap, vmin=-maxval, vmax=maxval)
506 im.set_clim(vmin=-maxval, vmax=maxval)
507 else:
508 tplarson 1.2 im.set_clim(vmin=0.0, vmax=maxval)
509 print("Scaling to maxval = %f"%maxval)
510 # else:
511 # im=ax.imshow(var, cmap=colormap)
|
512 tplarson 1.1
513 ax.set_axis_off()
514 fig.canvas.draw()
515
516 return (fig,im)
517
|
518 tplarson 1.2 def savefigure(ianimate=1, label=''):
|
519 tplarson 1.1
|
520 tplarson 1.3 l=lsave[-1]
521 m=msave[-1]
522 n=nsave[-1]
|
523 tplarson 1.2 if (ianimate == 0):
524 savestr='png_out/'+plotvar+label+'.l%im%in%i.png' % (l,m,n)
|
525 tplarson 1.1 plt.savefig(savestr, dpi=dpi)
526 print("File saved.")
527 else:
|
528 tplarson 1.2 pngdir = pngdirfmt.format(plotvar, label, l, m, n)
|
529 tplarson 1.1 if not os.path.exists(pngdir):
530 os.makedirs(pngdir)
531 print("Writing files...", end="")
532 for i in range(nframes):
|
533 tplarson 1.2 print(i, end=" ", flush=True)
|
534 tplarson 1.1 animate(i)
535 fpath = os.path.join(pngdir, '{:03d}.png'.format(i))
536 plt.savefig(fpath, dpi=dpi)
537 print("done.")
538
539
|
Return to
sosh.py
CVS log
Up to [Development] / JSOC / proj / globalhs / sosh