sdfr 200001.1.1__py3-none-win_amd64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
sdfr/sdf_helper.py ADDED
@@ -0,0 +1,2084 @@
1
+ import os
2
+ import re
3
+ try:
4
+ import numpy as np
5
+ import matplotlib.pyplot as plt
6
+ from matplotlib.transforms import Bbox
7
+ from matplotlib.offsetbox import HPacker, VPacker, TextArea, \
8
+ AnchoredOffsetbox
9
+ except:
10
+ pass
11
+ try:
12
+ from mpl_toolkits.axes_grid1 import make_axes_locatable
13
+ except:
14
+ try:
15
+ # Workaround for broken macOS installation
16
+ import sys
17
+ import matplotlib
18
+ sys.path.append(os.path.join(matplotlib.__path__[0],
19
+ '..', 'mpl_toolkits'))
20
+ from axes_grid1 import make_axes_locatable
21
+ except:
22
+ pass
23
+ try:
24
+ import builtins
25
+ except:
26
+ import __builtin__ as builtins
27
+
28
+ try:
29
+ import sdf
30
+ got_sdf = True
31
+ except ImportError:
32
+ got_sdf = False
33
+
34
+ try:
35
+ from matplotlib.pyplot import * # NOQA
36
+ got_mpl = True
37
+ except ImportError:
38
+ got_mpl = False
39
+
40
+
41
+ # mpl.rcParams['interactive'] = True
42
+ # hold()
43
+
44
+ global data, t, step, p, ppi, ppe, rho, ei, ee, vx, vy, vz, bx, by, bz
45
+ global x, y, z, xc, yc, zc, grid, grid_mid, mult_x, mult_y
46
+ global old_mtime, wkdir, verbose, fig, im, cbar
47
+
48
+ verbose = True
49
+
50
+ wkdir = 'Data'
51
+ old_mtime = 0
52
+ old_size = 0
53
+ old_filename = ''
54
+ cached = False
55
+ fig = None
56
+ im = None
57
+ cbar = None
58
+ mult_x = 1
59
+ mult_y = 1
60
+
61
+
62
+ class ic_type():
63
+ NEW = 1
64
+ RESTART = 2
65
+ SOD = 3
66
+ SALTZMAN = 4
67
+ NOH = 5
68
+ SEDOV = 6
69
+ BRIO = 7
70
+ WAVE = 8
71
+ ADVECT = 9
72
+
73
+
74
+ def get_si_prefix(scale, full_units=False):
75
+ scale = abs(scale)
76
+ mult = 1
77
+ sym = ''
78
+
79
+ if scale < 1e-24:
80
+ full_units = True
81
+ elif scale < 1e-21:
82
+ # yocto
83
+ mult = 1e24
84
+ sym = 'y'
85
+ elif scale < 1e-19:
86
+ # zepto
87
+ mult = 1e21
88
+ sym = 'z'
89
+ elif scale < 1e-16:
90
+ # atto
91
+ mult = 1e18
92
+ sym = 'a'
93
+ elif scale < 1e-13:
94
+ # femto
95
+ mult = 1e15
96
+ sym = 'f'
97
+ elif scale < 1e-10:
98
+ # pico
99
+ mult = 1e12
100
+ sym = 'p'
101
+ elif scale < 1e-7:
102
+ # nano
103
+ mult = 1e9
104
+ sym = 'n'
105
+ elif scale < 1e-4:
106
+ # micro
107
+ mult = 1e6
108
+ sym = '{\mu}'
109
+ elif scale < 1e-1:
110
+ # milli
111
+ mult = 1e3
112
+ sym = 'm'
113
+ elif scale >= 1e27:
114
+ full_units = True
115
+ elif scale >= 1e24:
116
+ # yotta
117
+ mult = 1e-24
118
+ sym = 'Y'
119
+ elif scale >= 1e21:
120
+ # zetta
121
+ mult = 1e-21
122
+ sym = 'Z'
123
+ elif scale >= 1e18:
124
+ # exa
125
+ mult = 1e-18
126
+ sym = 'E'
127
+ elif scale >= 1e15:
128
+ # peta
129
+ mult = 1e-15
130
+ sym = 'P'
131
+ elif scale >= 1e12:
132
+ # tera
133
+ mult = 1e-12
134
+ sym = 'T'
135
+ elif scale >= 1e9:
136
+ # giga
137
+ mult = 1e-9
138
+ sym = 'G'
139
+ elif scale >= 1e6:
140
+ # mega
141
+ mult = 1e-6
142
+ sym = 'M'
143
+ elif scale >= 1e3:
144
+ # kilo
145
+ mult = 1e-3
146
+ sym = 'k'
147
+
148
+ if full_units:
149
+ scale = scale * mult
150
+ if scale <= 0:
151
+ pwr = 0
152
+ else:
153
+ pwr = (-np.floor(np.log10(scale)))
154
+ mult = mult * np.power(10.0, pwr)
155
+ if np.rint(pwr) != 0:
156
+ sym = "(10^{%.0f})" % (-pwr) + sym
157
+
158
+ return mult, sym
159
+
160
+
161
+ def get_title(var=None, geom=False, extra_info=True):
162
+ global data
163
+
164
+ if var is not None:
165
+ if hasattr(var, 'blocklist'):
166
+ data = var.blocklist
167
+
168
+ t = data.Header['time']
169
+ mult, sym = get_si_prefix(t)
170
+
171
+ stitle = r'$t = {:.3}{}s$'.format(mult * t, sym)
172
+
173
+ if not extra_info:
174
+ return stitle
175
+
176
+ if hasattr(data, 'Logical_flags'):
177
+ if hasattr(data.Logical_flags, 'use_szp') \
178
+ and data.Logical_flags.use_szp:
179
+ stitle += r', $m_f = {:.1}$'.format(data.Real_flags.m_f)
180
+ if hasattr(data.Logical_flags, 'use_tts') \
181
+ and data.Logical_flags.use_tts:
182
+ stitle += r', TTS'
183
+ if hasattr(data.Logical_flags, 'use_tav') \
184
+ and data.Logical_flags.use_tav:
185
+ stitle += r', Tensor viscosity'
186
+ else:
187
+ if hasattr(data.Logical_flags, 'use_qmono') \
188
+ and data.Logical_flags.use_qmono:
189
+ stitle += r', Edge viscosity'
190
+ if hasattr(data.Logical_flags, 'use_edge') \
191
+ and data.Logical_flags.use_edge:
192
+ stitle += r', Edge viscosity'
193
+
194
+ if hasattr(data, 'Real_flags'):
195
+ if hasattr(data.Real_flags, 'visc1'):
196
+ stitle += r', $c_1 = ' + str(data.Real_flags.visc1) + '$'
197
+ if hasattr(data.Real_flags, 'visc2'):
198
+ stitle += r'$,\,c_2 = ' + str(data.Real_flags.visc2) + '$'
199
+
200
+ if geom:
201
+ if hasattr(data, 'Logical_flags'):
202
+ if hasattr(data.Logical_flags, 'use_rz') \
203
+ and data.Logical_flags.use_rz:
204
+ stitle += r', R-Z'
205
+ else:
206
+ stitle += r', Cartesian'
207
+
208
+ if hasattr(data.Logical_flags, 'polar_grid') \
209
+ and data.Logical_flags.polar_grid:
210
+ stitle += r' Polar'
211
+
212
+ return stitle
213
+
214
+
215
+ def get_default_iso(data):
216
+ iso = True
217
+ if hasattr(data, 'Integer_flags') \
218
+ and hasattr(data.Integer_flags, 'ic_type'):
219
+ ic = data.Integer_flags.ic_type
220
+ if ic == ic_type.NOH or ic == ic_type.SEDOV:
221
+ iso = True
222
+ return iso
223
+
224
+
225
+ def get_file_list(wkd=None, base=None, block=None):
226
+ """Get a list of SDF filenames containing sequence numbers
227
+
228
+ Parameters
229
+ ----------
230
+ wkd : str
231
+ The directory in which to search
232
+ If no other keyword arguments are passed, then the code will
233
+ automatically attempt to detect if this field is base or block
234
+ base : str
235
+ A representative filename or directory
236
+ block : sdf.Block or sdf.BlockList
237
+ A representative sdf dataset or block
238
+
239
+ Returns
240
+ -------
241
+ file_list : str array
242
+ An array of filenames
243
+ """
244
+ import os.path
245
+ import glob
246
+ global wkdir
247
+
248
+ if wkd is not None:
249
+ if os.path.isdir(wkd):
250
+ wkdir = wkd
251
+ elif os.path.exists(wkd):
252
+ base = wkd
253
+ elif isinstance(wkd, sdf.BlockList) \
254
+ or isinstance(wkd, sdf.Block) or type(wkd) is dict:
255
+ block = wkd
256
+
257
+ if base is None and block is not None:
258
+ if hasattr(block, 'blocklist'):
259
+ bl = block.blocklist
260
+ if hasattr(bl, 'Header') and 'filename' in bl.Header:
261
+ base = bl.Header['filename']
262
+ elif hasattr(block, 'Header') and 'filename' in block.Header:
263
+ base = block.Header['filename']
264
+
265
+ if base is not None:
266
+ if os.path.isfile(base[0]):
267
+ apath = os.path.abspath(base[0])
268
+ else:
269
+ apath = os.path.abspath(base)
270
+ wkdir = os.path.dirname(apath)
271
+ flist = glob.glob(wkdir + "/*.sdf")
272
+ flist.remove(apath)
273
+ flist = [apath] + sorted(flist)
274
+ else:
275
+ flist = glob.glob(wkdir + "/*[0-9][0-9]*.sdf")
276
+ if len(flist) == 0:
277
+ flist = glob.glob("*[0-9][0-9]*.sdf")
278
+ flist = sorted(flist)
279
+
280
+ return flist
281
+
282
+
283
+ def get_job_id(file_list=None, base=None, block=None):
284
+ """Get a representative job ID for a list of files
285
+
286
+ Parameters
287
+ ----------
288
+ file_list : str list
289
+ A list of filenames to search
290
+ If no other keyword arguments are passed, then the code will
291
+ automatically attempt to detect if this field is base or block
292
+ base : str
293
+ A representative filename or directory
294
+ block : sdf.Block or sdf.BlockList
295
+ A representative sdf dataset or block
296
+
297
+ Returns
298
+ -------
299
+ job_id : str
300
+ The job ID
301
+ """
302
+
303
+ if file_list is not None and type(file_list) is not list:
304
+ if os.path.exists(file_list):
305
+ base = file_list
306
+ file_list = None
307
+ elif isinstance(file_list, sdf.BlockList) \
308
+ or isinstance(file_list, sdf.Block) or type(file_list) is dict:
309
+ block = file_list
310
+ file_list = None
311
+
312
+ if block is not None and base is None:
313
+ if hasattr(block, 'blocklist'):
314
+ bl = block.blocklist
315
+ if hasattr(bl, 'Header') and 'filename' in bl.Header:
316
+ base = bl.Header['filename']
317
+ elif hasattr(block, 'Header') and 'filename' in block.Header:
318
+ base = block.Header['filename']
319
+
320
+ if base is not None:
321
+ try:
322
+ data = sdf.read(base, mmap=0)
323
+ if len(data.__dict__) > 1:
324
+ return data.Header['jobid1']
325
+ except:
326
+ pass
327
+
328
+ # Find the job id
329
+ if file_list is not None:
330
+ for f in file_list:
331
+ try:
332
+ data = sdf.read(f, mmap=0)
333
+ if len(data.__dict__) < 2:
334
+ continue
335
+ return data.Header['jobid1']
336
+ except:
337
+ pass
338
+
339
+ return None
340
+
341
+
342
+ def get_files(wkd=None, base=None, block=None, varname=None, fast=True):
343
+ """Get a list of SDF filenames belonging to the same run
344
+
345
+ Parameters
346
+ ----------
347
+ wkd : str
348
+ The directory in which to search
349
+ If no other keyword arguments are passed, then the code will
350
+ automatically attempt to detect if this field is base or block
351
+ base : str
352
+ A representative filename or directory
353
+ block : sdf.Block or sdf.BlockList
354
+ A representative sdf dataset or block
355
+ varname : str
356
+ A variable name that must be present in the file
357
+ fast : bool
358
+ Assume that files follow strict datestamp ordering and exit once
359
+ the first file that doesn't match the job ID
360
+
361
+ Returns
362
+ -------
363
+ file_list : str array
364
+ An array of filenames
365
+ """
366
+
367
+ if wkd is not None:
368
+ if os.path.isdir(wkd):
369
+ pass
370
+ elif os.path.exists(wkd):
371
+ base = wkd
372
+ elif isinstance(wkd, sdf.BlockList) \
373
+ or isinstance(wkd, sdf.Block) or type(wkd) is dict:
374
+ block = wkd
375
+
376
+ if block is not None and base is None:
377
+ if hasattr(block, 'blocklist'):
378
+ bl = block.blocklist
379
+ if hasattr(bl, 'Header') and 'filename' in bl.Header:
380
+ base = bl.Header['filename']
381
+ elif hasattr(block, 'Header') and 'filename' in block.Header:
382
+ base = block.Header['filename']
383
+
384
+ flist = get_file_list(wkd=wkd, base=base)
385
+ flist.sort(key=lambda x: os.path.getmtime(x))
386
+
387
+ job_id = get_job_id(flist, base=base, block=block)
388
+
389
+ # Add all files matching the job id
390
+ file_list = []
391
+ for f in reversed(flist):
392
+ try:
393
+ data = sdf.read(f, mmap=0, dict=True)
394
+ if len(data) < 2:
395
+ continue
396
+ file_job_id = data['Header']['jobid1']
397
+ if file_job_id == job_id:
398
+ if varname is None:
399
+ file_list.append(f)
400
+ elif varname in data:
401
+ file_list.append(f)
402
+ elif len(file_list) > 0:
403
+ break
404
+ except:
405
+ pass
406
+
407
+ return list(reversed(file_list))
408
+
409
+
410
+ def get_time(time=0, first=False, last=False, wkd=None, base=None, block=None,
411
+ fast=True):
412
+ """Get an SDF dataset that matches a given time
413
+
414
+ Parameters
415
+ ----------
416
+ time : float
417
+ The time to search for. If specified then the dateset that is
418
+ closest to this time will be returned
419
+ first : bool
420
+ If set to True then the dataset with the earliest simulation time
421
+ will be returned
422
+ last : bool
423
+ If set to True then the dataset with the latest simulation time
424
+ will be returned
425
+ wkd : str
426
+ The directory in which to search
427
+ base : str
428
+ A representative filename or directory
429
+ block : sdf.Block or sdf.BlockList
430
+ A representative sdf dataset or block
431
+ fast : bool
432
+ Use a faster but less thorough method for returning first/last
433
+
434
+ Returns
435
+ -------
436
+ data : sdf.BlockList
437
+ An SDF dataset
438
+ """
439
+ global data, wkdir
440
+
441
+ flist = get_file_list(wkd=wkd, base=base, block=block)
442
+
443
+ if len(flist) == 0:
444
+ print("No SDF files found")
445
+ return
446
+
447
+ flist.sort(key=lambda x: os.path.getmtime(x))
448
+ job_id = get_job_id(flist, base=base, block=block)
449
+
450
+ if time is None and not first:
451
+ last = True
452
+
453
+ t = None
454
+ fname = None
455
+ if last:
456
+ flist = list(reversed(flist))
457
+ t_old = -1e90
458
+ else:
459
+ t_old = 1e90
460
+
461
+ for f in flist:
462
+ dat_tmp = sdf.read(f)
463
+ if len(dat_tmp.__dict__) < 2:
464
+ continue
465
+ if job_id != dat_tmp.Header['jobid1']:
466
+ continue
467
+
468
+ t = dat_tmp.Header['time']
469
+ if last:
470
+ if fast:
471
+ fname = f
472
+ break
473
+ if t > t_old:
474
+ fname = f
475
+ t_old = t
476
+ elif first:
477
+ if fast:
478
+ fname = f
479
+ break
480
+ if t < t_old:
481
+ fname = f
482
+ t_old = t
483
+ else:
484
+ td = abs(t - time)
485
+ if td < t_old:
486
+ t_old = td
487
+ fname = f
488
+ if td < 1e-30:
489
+ # Exact match found. No need to search further
490
+ break
491
+
492
+ if fname is None:
493
+ raise Exception("No valid file found in directory: " + wkdir)
494
+
495
+ data = getdata(fname, verbose=False)
496
+ return data
497
+
498
+
499
+ def get_step(step=0, first=False, last=False, wkd=None, base=None, block=None,
500
+ fast=True):
501
+ """Get an SDF dataset that matches a given step
502
+
503
+ Parameters
504
+ ----------
505
+ step : int
506
+ The time to search for. If specified then the dateset that is
507
+ closest to this time will be returned
508
+ first : bool
509
+ If set to True then the dataset with the earliest simulation step
510
+ will be returned
511
+ last : bool
512
+ If set to True then the dataset with the latest simulation step
513
+ will be returned
514
+ wkd : str
515
+ The directory in which to search
516
+ base : str
517
+ A representative filename or directory
518
+ block : sdf.Block or sdf.BlockList
519
+ A representative sdf dataset or block
520
+ fast : bool
521
+ Use a faster but less thorough method for returning first/last
522
+
523
+ Returns
524
+ -------
525
+ data : sdf.BlockList
526
+ An SDF dataset
527
+ """
528
+ global data, wkdir
529
+
530
+ flist = get_file_list(wkd=wkd, base=base, block=block)
531
+
532
+ if len(flist) == 0:
533
+ print("No SDF files found")
534
+ return
535
+
536
+ flist.sort(key=lambda x: os.path.getmtime(x))
537
+ job_id = get_job_id(flist, base=base, block=block)
538
+
539
+ if step is None and not first:
540
+ last = True
541
+
542
+ t = None
543
+ fname = None
544
+ if last:
545
+ flist = list(reversed(flist))
546
+ t_old = -1e90
547
+ else:
548
+ t_old = 1e90
549
+
550
+ for f in flist:
551
+ dat_tmp = sdf.read(f)
552
+ if len(dat_tmp.__dict__) < 2:
553
+ continue
554
+ if job_id != dat_tmp.Header['jobid1']:
555
+ continue
556
+
557
+ t = dat_tmp.Header['step']
558
+ if last:
559
+ if fast:
560
+ fname = f
561
+ break
562
+ if t > t_old:
563
+ fname = f
564
+ t_old = t
565
+ elif first:
566
+ if fast:
567
+ fname = f
568
+ break
569
+ if t < t_old:
570
+ fname = f
571
+ t_old = t
572
+ else:
573
+ td = abs(t - step)
574
+ if td < t_old:
575
+ t_old = td
576
+ fname = f
577
+ if td == 0:
578
+ # Exact match found. No need to search further
579
+ break
580
+
581
+ if fname is None:
582
+ raise Exception("No valid file found in directory: " + wkdir)
583
+
584
+ data = getdata(fname, verbose=False)
585
+ return data
586
+
587
+
588
+ def get_latest(wkd=None, base=None, block=None):
589
+ """Get the latest SDF dataset in a directory
590
+
591
+ Parameters
592
+ ----------
593
+ wkd : str
594
+ The directory in which to search
595
+ If no other keyword arguments are passed, then the code will
596
+ automatically attempt to detect if this field is base or block
597
+ base : str
598
+ A representative filename or directory
599
+ block : sdf.Block or sdf.BlockList
600
+ A representative sdf dataset or block
601
+
602
+ Returns
603
+ -------
604
+ data : sdf.BlockList
605
+ An SDF dataset
606
+ """
607
+ return get_step(last=True, wkd=wkd, base=base, block=base)
608
+
609
+
610
+ def get_first(every=False, wkd=None, base=None, block=None):
611
+ if not every and not base:
612
+ base = get_newest_file(wkd=wkd, block=block)
613
+ return get_step(first=True, wkd=wkd, base=base, block=base)
614
+
615
+
616
+ def get_last(every=False, wkd=None, base=None, block=None):
617
+ if not every and not base:
618
+ base = get_newest_file(wkd=wkd, block=block)
619
+ return get_step(last=True, wkd=wkd, base=base, block=base)
620
+
621
+
622
+ def get_oldest_file(wkd=None, base=None, block=None):
623
+ import os.path
624
+
625
+ flist = get_file_list(wkd=wkd, base=base, block=block)
626
+ flist.sort(key=lambda x: os.path.getmtime(x))
627
+ for n in range(len(flist)):
628
+ f = flist[n]
629
+ data = sdf.read(f, mmap=0, dict=True)
630
+ if len(data) > 1:
631
+ return f
632
+
633
+ return None
634
+
635
+
636
+ def get_newest_file(wkd=None, base=None, block=None):
637
+ import os.path
638
+
639
+ flist = get_file_list(wkd=wkd, base=base, block=block)
640
+ flist.sort(key=lambda x: os.path.getmtime(x))
641
+ for n in range(len(flist)):
642
+ f = flist[-n-1]
643
+ data = sdf.read(f, mmap=0, dict=True)
644
+ if len(data) > 1:
645
+ return f
646
+
647
+ return None
648
+
649
+
650
+ def get_oldest(**kwargs):
651
+ return getdata(get_oldest_file(**kwargs), verbose=False)
652
+
653
+
654
+ def get_newest(**kwargs):
655
+ return getdata(get_newest_file(**kwargs), verbose=False)
656
+
657
+
658
+ def set_wkdir(wkd):
659
+ global wkdir
660
+ wkdir = wkd
661
+
662
+
663
+ def get_wkdir():
664
+ global wkdir
665
+ return wkdir
666
+
667
+
668
+ def sdfr(filename):
669
+ return sdf.read(filename)
670
+
671
+
672
+ def plot_auto(*args, **kwargs):
673
+ try:
674
+ dims = args[0].dims
675
+ except:
676
+ print('error: Variable cannot be auto determined. '
677
+ + 'Use plot1d or plot2d')
678
+ return
679
+ if (len(dims) == 1):
680
+ if type(args[0]) is sdf.BlockStitchedPath:
681
+ plot_rays(*args, **kwargs)
682
+ elif (len(args[0].grid.dims) == 1):
683
+ plot1d(*args, **kwargs)
684
+ else:
685
+ plot_path(*args, **kwargs)
686
+ elif (len(dims) == 2):
687
+ k = 'set_ylabel'
688
+ if k in kwargs:
689
+ del kwargs[k]
690
+ plot2d(*args, **kwargs)
691
+ elif len(dims) == 3 and \
692
+ ('ix' in kwargs or 'iy' in kwargs or 'iz' in kwargs):
693
+ k = 'set_ylabel'
694
+ if k in kwargs:
695
+ del kwargs[k]
696
+ plot2d(*args, **kwargs)
697
+ else:
698
+ print('error: Unable to plot variables of this dimensionality')
699
+
700
+
701
+ def oplot_auto(*args, **kwargs):
702
+ kwargs['set_ylabel'] = False
703
+ kwargs['hold'] = True
704
+ plot_auto(*args, **kwargs)
705
+
706
+
707
+ def oplot1d(*args, **kwargs):
708
+ kwargs['set_ylabel'] = False
709
+ kwargs['hold'] = True
710
+ plot1d(*args, **kwargs)
711
+
712
+
713
+ def plot1d(var, fmt=None, xdir=None, idx=-1, xscale=0, yscale=0, scale=0,
714
+ cgs=False, title=True, sym=True, set_ylabel=True, hold=True,
715
+ subplot=None, figure=None, **kwargs):
716
+ global data
717
+ global x, y, mult_x, mult_y
718
+
719
+ if len(var.dims) != 1 and len(var.dims) != 2:
720
+ print("error: Not a 1d dataset")
721
+ return
722
+
723
+ if subplot is None:
724
+ if figure is None:
725
+ figure = plt.gcf()
726
+ if len(figure.get_axes()) == 0:
727
+ hold = False
728
+ # Only clear the figure if one isn't supplied by the user
729
+ if not hold:
730
+ try:
731
+ figure.clf()
732
+ except:
733
+ pass
734
+ elif not hold:
735
+ figure.clf()
736
+ # Have to add subplot after clearing figure
737
+ subplot = plt.subplot(111)
738
+ elif figure is None:
739
+ figure = subplot.figure
740
+ if not hold:
741
+ subplot.clear()
742
+ if hasattr(subplot, 'colorbar'):
743
+ subplot.colorbar.remove()
744
+
745
+ if var.dims[0] == var.grid.dims[0]:
746
+ grid = var.grid
747
+ else:
748
+ grid = var.grid_mid
749
+
750
+ if len(var.dims) > 1:
751
+ if xdir is None:
752
+ if var.dims[1] < var.dims[0]:
753
+ xdir = 0
754
+ else:
755
+ xdir = 1
756
+ if xdir == 0:
757
+ if idx == -1:
758
+ idx = int(var.dims[1] / 2)
759
+ s = (slice(None), idx)
760
+ else:
761
+ if idx == -1:
762
+ idx = int(var.dims[0] / 2)
763
+ s = (idx, slice(None))
764
+ Y = var.data[s]
765
+ else:
766
+ xdir = 0
767
+ Y = var.data
768
+
769
+ if np.ndim(var.grid.data[0]) == 1:
770
+ X = grid.data[xdir]
771
+ else:
772
+ X = grid.data[xdir][s]
773
+
774
+ if xdir is None:
775
+ xdir = 0
776
+
777
+ if scale > 0:
778
+ xscale = scale
779
+ yscale = scale
780
+
781
+ # This is an attempt to get both axes to use the same scale
782
+ # I'm not at all confident that it always does the right thing
783
+ if xscale == 0 and yscale == 0:
784
+ xlen = max(abs(X[0]), abs(X[-1]))
785
+ ylen = max(abs(Y[0]), abs(Y[-1]))
786
+ ratio1 = ylen / xlen
787
+ ratio2 = xlen / ylen
788
+ if ratio1 >= 1 and ratio1 < 1e3:
789
+ xscale = xlen
790
+ yscale = xlen
791
+ elif ratio2 >= 1 and ratio2 < 1e3:
792
+ xscale = ylen
793
+ yscale = ylen
794
+
795
+ if xscale == 0:
796
+ length = max(abs(X[0]), abs(X[-1]))
797
+ mult_x, sym_x = get_si_prefix(length)
798
+ else:
799
+ mult_x, sym_x = get_si_prefix(xscale)
800
+
801
+ if yscale == 0:
802
+ length = max(abs(Y[0]), abs(Y[-1]))
803
+ mult_y, sym_y = get_si_prefix(length)
804
+ else:
805
+ mult_y, sym_y = get_si_prefix(yscale)
806
+
807
+ X = mult_x * X
808
+ Y = mult_y * Y
809
+
810
+ if fmt:
811
+ subplot.plot(X, Y, fmt, **kwargs)
812
+ else:
813
+ subplot.plot(X, Y, **kwargs)
814
+
815
+ subplot.set_xlabel(grid.labels[xdir] + ' $('
816
+ + escape_latex(sym_x + grid.units[xdir]) + ')$')
817
+ if set_ylabel:
818
+ subplot.set_ylabel(var.name + ' $('
819
+ + escape_latex(sym_y + var.units) + ')$')
820
+
821
+ if title:
822
+ subplot.set_title(get_title(var), fontsize='large', y=1.03)
823
+
824
+ figure.set_tight_layout(True)
825
+ figure.canvas.draw()
826
+
827
+
828
+ def plot_path(var, xdir=None, ydir=None, xscale=0, yscale=0, scale=0,
829
+ title=True, hold=True, subplot=None, figure=None, iso=True,
830
+ add_cbar=True, cbar_label=True, cbar_wd=5, cbar_top=False,
831
+ svar=None, update=True, axis_only=False, clip_reflect=False,
832
+ power=(-3,3), **kwargs):
833
+ """Plot an SDF path variable (eg. a laser ray)
834
+
835
+ Parameters
836
+ ----------
837
+ var : sdf.Block
838
+ The SDF block for the path to plot
839
+ xdir : integer
840
+ The dimension to use for the x-direction. If the path is 2D then
841
+ ydir will automatically be selected.
842
+ ydir : integer
843
+ The dimension to use for the y-direction. If the path is 2D then
844
+ xdir will automatically be selected.
845
+ xscale : real
846
+ Value to use for scaling the x-axis. If not set then the x-axis will
847
+ be scaled automatically. Set this to 1 to disable scaling.
848
+ yscale : real
849
+ Value to use for scaling the y-axis. If not set then the x-axis will
850
+ be scaled automatically. Set this to 1 to disable scaling.
851
+ scale : real
852
+ Value to use for scaling both axes. If not set then the axes will
853
+ be scaled automatically. Set this to 1 to disable scaling.
854
+ title : logical or string
855
+ If set to False, don't add a title to the plot.
856
+ Otherwise, add the specified value if it is a string or
857
+ automatically generate one if it is just True.
858
+ hold : logical
859
+ If True, do not clear the figure before plotting.
860
+ subplot :
861
+ The subplot to use for plotting
862
+ figure :
863
+ The figure to use for plotting
864
+ iso : logical
865
+ If True, generate an isotropic plot (x/y axis equal).
866
+ add_cbar : logical
867
+ Add a colorbar if true
868
+ cbar_label : logical or string
869
+ If set to False, don't add a label to the colorbar.
870
+ Otherwise, add the specified value if it is a string or
871
+ automatically generate one if it is just True.
872
+ cbar_wd : integer
873
+ The width of the colorbar
874
+ cbar_top : logical
875
+ If set to true, the colorbar is plotted along the top of the figure
876
+ instead of the right-hand side
877
+ update : logical
878
+ If set to true then update the axis limits for this path
879
+ axis_only : logical
880
+ If set to true then only update the axis limits for this path
881
+ svar : sdf.Block
882
+ If set, use the extents of this variable to set the axis range for
883
+ this plot
884
+ clip_reflect : logical
885
+ If set to true, then rays are clipped at the point where the path
886
+ gradient is either zero or huge
887
+ power : sequence
888
+ Sets size thresholds for scientific notation.
889
+ Two-element sequence passed throuth to set_powerlimits()
890
+
891
+ **kwargs : dict
892
+ All other keyword arguments are passed to matplotlib plotting
893
+ routine
894
+ """
895
+ import matplotlib.pyplot as plt
896
+ from matplotlib.collections import LineCollection
897
+ global data
898
+ global x, y, mult_x, mult_y
899
+
900
+ if 'norm_values' not in plot_path.__dict__:
901
+ plot_path.norm_values = None
902
+ plot_path.axis = None
903
+
904
+ if axis_only:
905
+ if plot_path.axis is None:
906
+ return
907
+ if subplot is None:
908
+ if figure is None:
909
+ figure = plt.gcf()
910
+ subplot = plt.subplot(111)
911
+ elif figure is None:
912
+ figure = subplot.figure
913
+ subplot.axis(plot_path.axis)
914
+ figure.set_tight_layout(True)
915
+ figure.canvas.draw()
916
+ return
917
+
918
+ if len(var.dims) != 1:
919
+ print("error: Not a 1d dataset")
920
+ return
921
+
922
+ if len(plt.get_fignums()) == 0:
923
+ hold = False
924
+
925
+ if subplot is None:
926
+ if figure is None:
927
+ figure = plt.gcf()
928
+ if len(figure.get_axes()) == 0:
929
+ hold = False
930
+ # Only clear the figure if one isn't supplied by the user
931
+ if not hold:
932
+ try:
933
+ figure.clf()
934
+ except:
935
+ pass
936
+ elif not hold:
937
+ figure.clf()
938
+ # Have to add subplot after clearing figure
939
+ subplot = plt.subplot(111)
940
+ elif figure is None:
941
+ figure = subplot.figure
942
+ if not hold:
943
+ subplot.clear()
944
+ if hasattr(subplot, 'colorbar'):
945
+ subplot.colorbar.remove()
946
+
947
+ if not hold:
948
+ plot_path.norm_values = None
949
+ plot_path.axis = None
950
+
951
+ #if var.dims[0] == var.grid.dims[0]:
952
+ # grid = var.grid
953
+ #else:
954
+ # grid = var.grid_mid
955
+ grid = var.grid
956
+
957
+ test_dir = False
958
+ if xdir is None:
959
+ xdir = 0
960
+ else:
961
+ test_dir = True
962
+
963
+ if ydir is None:
964
+ ydir = 1
965
+ else:
966
+ test_dir = True
967
+
968
+ if test_dir:
969
+ dimensions = np.shape(grid.data)[0]
970
+ xdir = xdir % dimensions
971
+ ydir = ydir % dimensions
972
+ if xdir == ydir:
973
+ ydir = (xdir + 1) % dimensions
974
+
975
+ X = grid.data[xdir]
976
+ Y = grid.data[ydir]
977
+
978
+ if not hold:
979
+ if scale > 0:
980
+ xscale = scale
981
+ yscale = scale
982
+
983
+ # This is an attempt to get both axes to use the same scale
984
+ # I'm not at all confident that it always does the right thing
985
+ if xscale == 0 and yscale == 0:
986
+ xlen = max(abs(X[0]), abs(X[-1]))
987
+ ylen = max(abs(Y[0]), abs(Y[-1]))
988
+ ratio1 = ylen / xlen
989
+ ratio2 = xlen / ylen
990
+ if ratio1 >= 1 and ratio1 < 1e3:
991
+ xscale = xlen
992
+ yscale = xlen
993
+ elif ratio2 >= 1 and ratio2 < 1e3:
994
+ xscale = ylen
995
+ yscale = ylen
996
+
997
+ if xscale == 0:
998
+ length = max(abs(X[0]), abs(X[-1]))
999
+ mult_x, sym_x = get_si_prefix(length)
1000
+ else:
1001
+ mult_x, sym_x = get_si_prefix(xscale)
1002
+
1003
+ if yscale == 0:
1004
+ length = max(abs(Y[0]), abs(Y[-1]))
1005
+ mult_y, sym_y = get_si_prefix(length)
1006
+ else:
1007
+ mult_y, sym_y = get_si_prefix(yscale)
1008
+
1009
+ X = mult_x * X
1010
+ Y = mult_y * Y
1011
+ c = var.data
1012
+
1013
+ if clip_reflect:
1014
+ g = (X[0] - X[1]) / (Y[0] - Y[1])
1015
+ for i in range(len(c) - 1):
1016
+ ddx = X[i] - X[i+1]
1017
+ ddy = Y[i] - Y[i+1]
1018
+ if abs(ddy) < 1e-16:
1019
+ grad = 100
1020
+ else:
1021
+ grad = ddx / ddy / g
1022
+ if abs(grad) > 10 or grad < 0:
1023
+ X = np.copy(X[:i+1])
1024
+ Y = np.copy(Y[:i+1])
1025
+ c = np.copy(c[:i+1])
1026
+ break
1027
+
1028
+ points = np.array([X, Y]).T.reshape(-1, 1, 2)
1029
+ segments = np.concatenate([points[:-1], points[1:]], axis=1)
1030
+
1031
+ if not hold or plot_path.norm_values is None:
1032
+ k = 'vmin'
1033
+ k1 = 'vrange'
1034
+ if k in kwargs:
1035
+ vmin = kwargs[k]
1036
+ elif k1 in kwargs:
1037
+ vmin = kwargs[k1][0]
1038
+ else:
1039
+ vmin = c.min()
1040
+
1041
+ k = 'vmax'
1042
+ if k in kwargs:
1043
+ vmax = kwargs[k]
1044
+ elif k1 in kwargs:
1045
+ vmax = kwargs[k1][1]
1046
+ else:
1047
+ vmax = c.max()
1048
+
1049
+ plot_path.norm_values = plt.Normalize(vmin, vmax)
1050
+
1051
+ kk = {}
1052
+ k = 'lw'
1053
+ if k in kwargs:
1054
+ kk[k] = kwargs[k]
1055
+ k = 'linewidth'
1056
+ if k in kwargs:
1057
+ kk[k] = kwargs[k]
1058
+ lc = LineCollection(segments, norm=plot_path.norm_values, **kk)
1059
+ lc.set_array(c)
1060
+ im = subplot.add_collection(lc)
1061
+
1062
+ if not hold:
1063
+ subplot.set_xlabel(grid.labels[xdir] + ' $('
1064
+ + escape_latex(sym_x + grid.units[xdir]) + ')$')
1065
+ subplot.set_ylabel(grid.labels[ydir] + ' $('
1066
+ + escape_latex(sym_y + grid.units[ydir]) + ')$')
1067
+
1068
+ if title:
1069
+ if type(title) is str:
1070
+ title_label = title
1071
+ else:
1072
+ title_label = get_title(var, extra_info=False)
1073
+ subplot.set_title(title_label, fontsize='large', y=1.03)
1074
+
1075
+ subplot.axis('tight')
1076
+ if iso:
1077
+ subplot.axis('image')
1078
+
1079
+ plot_path.axis = subplot.axis()
1080
+ if update:
1081
+ subplot.axis([X.min(), X.max(), Y.min(), Y.max()])
1082
+ plot_path.axis = [X.min(), X.max(), Y.min(), Y.max()]
1083
+
1084
+ if not hold and add_cbar:
1085
+ ax = subplot.axes
1086
+ ca = subplot
1087
+ divider = make_axes_locatable(ca)
1088
+ pad = int(0.6 * cbar_wd + 0.5)
1089
+ if cbar_top:
1090
+ cax = divider.append_axes("top", "%i%%" % cbar_wd,
1091
+ pad="%i%%" % pad)
1092
+ cbar = figure.colorbar(im, orientation='horizontal',
1093
+ cax=cax, ax=ax)
1094
+ cax.xaxis.set_ticks_position('top')
1095
+ else:
1096
+ cax = divider.append_axes("right", "%i%%" % cbar_wd,
1097
+ pad="%i%%" % pad)
1098
+ cbar = figure.colorbar(im, cax=cax, ax=ax)
1099
+ try:
1100
+ cbar.formatter.set_powerlimits(power)
1101
+ cbar.update_ticks()
1102
+ except:
1103
+ pass
1104
+ subplot.colorbar = cax
1105
+ plt.sca(ax)
1106
+ if cbar_label:
1107
+ if type(cbar_label) is str:
1108
+ var_label = cbar_label
1109
+ else:
1110
+ var_label = var.name + ' $(' + escape_latex(var.units) + ')$'
1111
+ if cbar_top:
1112
+ cbar.set_label(var_label, fontsize='large')
1113
+ cax.xaxis.set_label_position('top')
1114
+ else:
1115
+ cbar.set_label(var_label, fontsize='large', x=1.2)
1116
+
1117
+ if not hold and svar is not None:
1118
+ lim = np.array(svar.grid.extents)
1119
+ lim[0] *= mult_x
1120
+ lim[1] *= mult_x
1121
+ lim[2] *= mult_y
1122
+ lim[3] *= mult_y
1123
+ subplot.axis([lim[0], lim[2], lim[1], lim[3]])
1124
+ elif hold:
1125
+ if plot_path.axis is None:
1126
+ plot_path.axis = [X.min(), X.max(), Y.min(), Y.max()]
1127
+ lims = plot_path.axis
1128
+ if lims is None:
1129
+ lim = [X.min(), X.max(), Y.min(), Y.max()]
1130
+ else:
1131
+ lim = list(lims)
1132
+ v = X.min()
1133
+ if v < lim[0]:
1134
+ lim[0] = v
1135
+ v = X.max()
1136
+ if v > lim[1]:
1137
+ lim[1] = v
1138
+ v = Y.min()
1139
+ if v < lim[2]:
1140
+ lim[2] = v
1141
+ v = Y.max()
1142
+ if v > lim[3]:
1143
+ lim[3] = v
1144
+ plot_path.axis = lim
1145
+ if update:
1146
+ subplot.axis(lim)
1147
+
1148
+ if update:
1149
+ figure.set_tight_layout(True)
1150
+ figure.canvas.draw()
1151
+
1152
+
1153
+ def plot_rays(var, skip=1, rays=None, **kwargs):
1154
+ """Plot all rays found in an SDF file
1155
+
1156
+ Parameters
1157
+ ----------
1158
+ var : sdf.Block
1159
+ The SDF variable for the rays to plot
1160
+ skip : integer
1161
+ Number of rays to skip before selecting the next one to plot
1162
+ """
1163
+
1164
+ ray_start = -1
1165
+ l = 'ray_start'
1166
+ if l in kwargs:
1167
+ ray_start = kwargs[l]
1168
+
1169
+ ray_stop = 1e9
1170
+ l = 'ray_stop'
1171
+ if l in kwargs:
1172
+ ray_stop = kwargs[l]
1173
+
1174
+ if type(var) is sdf.BlockStitchedPath:
1175
+ v = var.data[0]
1176
+ l = '_label'
1177
+ if l not in kwargs:
1178
+ kwargs[l] = var.name
1179
+
1180
+ if type(v) is sdf.BlockStitchedPath:
1181
+ for v in var.data:
1182
+ plot_rays(v, skip=skip, rays=rays, **kwargs)
1183
+ kwargs['hold'] = True
1184
+ return
1185
+
1186
+ k = 'cbar_label'
1187
+ if k not in kwargs or (kwargs[k] and type(kwargs[k]) != str):
1188
+ kwargs[k] = kwargs[l] + ' $(' + escape_latex(v.units) + ')$'
1189
+ del kwargs[l]
1190
+
1191
+ k0 = 'vmin'
1192
+ k1 = 'vmax'
1193
+ k = 'vrange'
1194
+ if k not in kwargs and not (k0 in kwargs and k1 in kwargs):
1195
+ v = var.data[0]
1196
+ vmin = v.data.min()
1197
+ vmax = v.data.max()
1198
+ for iray, v in enumerate(var.data):
1199
+ if iray < ray_start:
1200
+ continue
1201
+ if iray > ray_stop:
1202
+ break
1203
+ if iray%skip == 0:
1204
+ vmin = min(vmin, v.data.min())
1205
+ vmax = max(vmax, v.data.max())
1206
+ if k0 not in kwargs:
1207
+ kwargs[k0] = vmin
1208
+ if k1 not in kwargs:
1209
+ kwargs[k1] = vmax
1210
+
1211
+ for iray, v in enumerate(var.data):
1212
+ if iray < ray_start:
1213
+ continue
1214
+ if iray > ray_stop:
1215
+ break
1216
+ if iray%skip == 0:
1217
+ plot_auto(v, update=False, **kwargs)
1218
+ kwargs['hold'] = True
1219
+
1220
+ plot_auto(var.data[0], axis_only=True, **kwargs)
1221
+ kwargs['hold'] = True
1222
+
1223
+ return
1224
+
1225
+ split_name = var.name.split('/')
1226
+ start = split_name[0] + '_'
1227
+ end = '_' + split_name[-1]
1228
+ data = var.blocklist.__dict__
1229
+
1230
+ k0 = 'vmin'
1231
+ k1 = 'vmax'
1232
+ k = 'vrange'
1233
+ if k not in kwargs and not (k0 in kwargs and k1 in kwargs):
1234
+ vmin = var.data.min()
1235
+ vmax = var.data.max()
1236
+ iray = -1
1237
+ for k in data.keys():
1238
+ if k.startswith(start) and k.endswith(end):
1239
+ iray += 1
1240
+ if iray < ray_start:
1241
+ continue
1242
+ if iray > ray_stop:
1243
+ break
1244
+ if iray%skip == 0:
1245
+ vmin = min(vmin, data[k].data.min())
1246
+ vmax = max(vmax, data[k].data.max())
1247
+ if k0 not in kwargs:
1248
+ kwargs[k0] = vmin
1249
+ if k1 not in kwargs:
1250
+ kwargs[k1] = vmax
1251
+
1252
+ k = 'cbar_label'
1253
+ if k not in kwargs or (kwargs[k] and type(kwargs[k]) != str):
1254
+ # Remove /Ray[1-9]+ from the name
1255
+ kwargs[k] = '/'.join([split_name[0]] + split_name[2:]) \
1256
+ + ' $(' + escape_latex(var.units) + ')$'
1257
+
1258
+ iray = -1
1259
+ for k in data.keys():
1260
+ if k.startswith(start) and k.endswith(end):
1261
+ iray += 1
1262
+ if iray < ray_start:
1263
+ continue
1264
+ if iray > ray_stop:
1265
+ break
1266
+ if iray%skip == 0:
1267
+ plot_auto(data[k], hold=True, update=False, **kwargs)
1268
+ kwargs['hold'] = True
1269
+
1270
+ plot_auto(var, axis_only=True, **kwargs)
1271
+ kwargs['hold'] = True
1272
+
1273
+
1274
+ def oplot2d(*args, **kwargs):
1275
+ kwargs['hold'] = True
1276
+ plot2d(*args, **kwargs)
1277
+
1278
+
1279
+ def plot2d_array(array, x, y, extents, var_label, xlabel, ylabel, idx=None,
1280
+ iso=None, fast=None, title=True, full=True, vrange=None,
1281
+ reflect=0, norm=None, hold=True, xscale=0, yscale=0, scale=0,
1282
+ figure=None, subplot=None, add_cbar=True, cbar_label=True,
1283
+ cbar_wd=5, cbar_top=False, power=(-3,3), **kwargs):
1284
+ import matplotlib as mpl
1285
+ global data, fig, im, cbar
1286
+ global mult_x, mult_y
1287
+
1288
+ if idx is None:
1289
+ i0 = 0
1290
+ i1 = 1
1291
+ i2 = 2
1292
+ i3 = 3
1293
+ else:
1294
+ i0 = idx[0]
1295
+ i1 = idx[1]
1296
+ i2 = idx[2]
1297
+ i3 = idx[3]
1298
+
1299
+ cmap = plt.get_cmap()
1300
+
1301
+ if norm is not None:
1302
+ v0 = np.min(array) - norm
1303
+ v1 = np.max(array) - norm
1304
+
1305
+ if abs(v0/v1) > 1:
1306
+ low = 0
1307
+ high = 0.5 * (1 - v1/v0)
1308
+ else:
1309
+ low = 0.5 * (1 + v0/v1)
1310
+ high = 1.0
1311
+
1312
+ cmap = mpl.colors.LinearSegmentedColormap.from_list(
1313
+ 'tr', cmap(np.linspace(low, high, 256)))
1314
+
1315
+ if subplot is None:
1316
+ if figure is None:
1317
+ figure = plt.gcf()
1318
+ if len(figure.get_axes()) == 0:
1319
+ hold = False
1320
+ # Only clear the figure if one isn't supplied by the user
1321
+ if not hold:
1322
+ try:
1323
+ figure.clf()
1324
+ except:
1325
+ pass
1326
+ elif not hold:
1327
+ figure.clf()
1328
+ # Have to add subplot after clearing figure
1329
+ subplot = plt.subplot(111)
1330
+ elif figure is None:
1331
+ figure = subplot.figure
1332
+ if not hold:
1333
+ subplot.clear()
1334
+ if hasattr(subplot, 'colorbar'):
1335
+ subplot.colorbar.remove()
1336
+
1337
+ if iso is None:
1338
+ iso = get_default_iso(data)
1339
+
1340
+ if scale > 0:
1341
+ xscale = scale
1342
+ yscale = scale
1343
+
1344
+ ext = extents[:]
1345
+ # This is an attempt to get both axes to use the same scale
1346
+ # I'm not at all confident that it always does the right thing
1347
+ if xscale == 0 and yscale == 0:
1348
+ xlen = max(abs(extents[i2]), abs(extents[i0]))
1349
+ ylen = max(abs(extents[i3]), abs(extents[i1]))
1350
+ ratio1 = ylen / xlen
1351
+ ratio2 = xlen / ylen
1352
+ if ratio1 >= 1 and ratio1 < 1e3:
1353
+ xscale = xlen
1354
+ yscale = xlen
1355
+ elif ratio2 >= 1 and ratio2 < 1e3:
1356
+ xscale = ylen
1357
+ yscale = ylen
1358
+
1359
+ if xscale == 0:
1360
+ length = max(abs(extents[i2]), abs(extents[i0]))
1361
+ mult_x, sym_x = get_si_prefix(length)
1362
+ else:
1363
+ mult_x, sym_x = get_si_prefix(xscale)
1364
+
1365
+ if yscale == 0:
1366
+ length = max(abs(extents[i3]), abs(extents[i1]))
1367
+ mult_y, sym_y = get_si_prefix(length)
1368
+ else:
1369
+ mult_y, sym_y = get_si_prefix(yscale)
1370
+
1371
+ if vrange == 1:
1372
+ v = np.max(abs(array))
1373
+ vrange = [-v, v]
1374
+
1375
+ if fast:
1376
+ if reflect == 1:
1377
+ # about x=0
1378
+ ext[i0] = 2 * extents[i0] - extents[i2]
1379
+ array = np.vstack((np.flipud(array), array))
1380
+ elif reflect == 2:
1381
+ # about y=0
1382
+ ext[i1] = 2 * extents[i1] - extents[i3]
1383
+ array = np.hstack((np.fliplr(array), array))
1384
+ elif reflect == 3:
1385
+ # about x=0, y=0
1386
+ ext[i0] = 2 * extents[i0] - extents[i2]
1387
+ ext[i1] = 2 * extents[i1] - extents[i3]
1388
+ array = np.vstack((np.flipud(array), array))
1389
+ array = np.hstack((np.fliplr(array), array))
1390
+
1391
+ if np.ndim(x) == 1:
1392
+ if fast is None:
1393
+ fast = True
1394
+
1395
+ if not fast:
1396
+ Y, X = np.meshgrid(y, x)
1397
+ else:
1398
+ if fast is None:
1399
+ fast = False
1400
+
1401
+ if not fast:
1402
+ X = x
1403
+ Y = y
1404
+
1405
+ if fast:
1406
+ ext[i0] = mult_x * ext[i0]
1407
+ ext[i1] = mult_y * ext[i1]
1408
+ ext[i2] = mult_x * ext[i2]
1409
+ ext[i3] = mult_y * ext[i3]
1410
+ e = ext[i1]
1411
+ ext[i1] = ext[i2]
1412
+ ext[i2] = e
1413
+ ext = [ext[i0], ext[i1], ext[i2], ext[i3]]
1414
+ k = 'interpolation'
1415
+ if k not in kwargs:
1416
+ kwargs[k] = 'none'
1417
+ if vrange is None:
1418
+ im = subplot.imshow(array.T, origin='lower', extent=ext, cmap=cmap,
1419
+ **kwargs)
1420
+ else:
1421
+ im = subplot.imshow(array.T, origin='lower', extent=ext, cmap=cmap,
1422
+ vmin=vrange[0], vmax=vrange[1], **kwargs)
1423
+ else:
1424
+ X = np.multiply(mult_x, X)
1425
+ Y = np.multiply(mult_y, Y)
1426
+ if vrange is None:
1427
+ im = subplot.pcolormesh(X, Y, array, cmap=cmap, **kwargs)
1428
+ else:
1429
+ im = subplot.pcolormesh(X, Y, array, cmap=cmap,
1430
+ vmin=vrange[0], vmax=vrange[1], **kwargs)
1431
+
1432
+ subplot.set_xlabel(xlabel)
1433
+ subplot.set_ylabel(ylabel)
1434
+
1435
+ title_text = None
1436
+ if full:
1437
+ if add_cbar and cbar_label:
1438
+ title_text = get_title()
1439
+ else:
1440
+ title_text = var_label + ', ' + get_title()
1441
+ elif title:
1442
+ if not (add_cbar and cbar_label):
1443
+ title_text = var_label
1444
+
1445
+ if title and title_text:
1446
+ subplot.set_title(title_text, fontsize='large', y=1.03)
1447
+
1448
+ subplot.axis('tight')
1449
+ if iso:
1450
+ subplot.axis('image')
1451
+
1452
+ if not hold and add_cbar:
1453
+ ax = subplot.axes
1454
+ ca = subplot
1455
+ divider = make_axes_locatable(ca)
1456
+ pad = int(0.6 * cbar_wd + 0.5)
1457
+ if cbar_top:
1458
+ cax = divider.append_axes("top", "%i%%" % cbar_wd,
1459
+ pad="%i%%" % pad)
1460
+ cbar = figure.colorbar(im, orientation='horizontal',
1461
+ cax=cax, ax=ax)
1462
+ cax.xaxis.set_ticks_position('top')
1463
+ else:
1464
+ cax = divider.append_axes("right", "%i%%" % cbar_wd,
1465
+ pad="%i%%" % pad)
1466
+ cbar = figure.colorbar(im, cax=cax, ax=ax)
1467
+ try:
1468
+ cbar.formatter.set_powerlimits(power)
1469
+ cbar.update_ticks()
1470
+ except:
1471
+ pass
1472
+ subplot.colorbar = cax
1473
+ plt.sca(ax)
1474
+ if cbar_label:
1475
+ if type(cbar_label) is str:
1476
+ var_label = cbar_label
1477
+ if cbar_top:
1478
+ cbar.set_label(var_label, fontsize='large')
1479
+ cax.xaxis.set_label_position('top')
1480
+ else:
1481
+ cbar.set_label(var_label, fontsize='large', x=1.2)
1482
+ figure.canvas.draw()
1483
+
1484
+ figure.set_tight_layout(True)
1485
+ figure.canvas.draw()
1486
+
1487
+
1488
+ def plot2d(var, iso=None, fast=None, title=True, full=True, vrange=None,
1489
+ ix=None, iy=None, iz=None, reflect=0, norm=None, irange=None,
1490
+ jrange=None, hold=True, xscale=0, yscale=0, scale=0, figure=None,
1491
+ subplot=None, add_cbar=True, cbar_label=True, cbar_top=False,
1492
+ **kwargs):
1493
+ global data, fig, im, cbar
1494
+ global x, y, mult_x, mult_y
1495
+
1496
+ if type(irange) is list or type(irange) is tuple:
1497
+ si = slice(*irange)
1498
+ else:
1499
+ si = slice(None, irange)
1500
+
1501
+ if type(jrange) is list or type(jrange) is tuple:
1502
+ sj = slice(*jrange)
1503
+ else:
1504
+ sj = slice(None, jrange)
1505
+
1506
+ i0 = 0
1507
+ i1 = 1
1508
+ if len(var.dims) == 3:
1509
+ if ix is not None:
1510
+ if iz is None:
1511
+ if ix < 0:
1512
+ ix = int(var.dims[0] / 2)
1513
+ i0 = 1
1514
+ i1 = 2
1515
+ ss = ix, si, sj
1516
+ else:
1517
+ if ix < 0:
1518
+ ix = int(var.dims[2] / 2)
1519
+ i0 = 0
1520
+ i1 = 2
1521
+ ss = si, sj, ix
1522
+ elif iy is not None:
1523
+ if iy < 0:
1524
+ iy = int(var.dims[1] / 2)
1525
+ i0 = 0
1526
+ i1 = 2
1527
+ ss = si, iy, sj
1528
+ if iz is not None:
1529
+ i0 = 0
1530
+ i1 = 1
1531
+ ss = si, iy, sj
1532
+ elif iz is not None:
1533
+ if iz < 0:
1534
+ iz = int(var.dims[2] / 2)
1535
+ i0 = 0
1536
+ i1 = 1
1537
+ ss = si, sj, iz
1538
+ else:
1539
+ print("error: Not a 2d dataset")
1540
+ return
1541
+ i2 = i0 + 3
1542
+ i3 = i1 + 3
1543
+ elif len(var.dims) != 2:
1544
+ print("error: Not a 2d dataset")
1545
+ return
1546
+ else:
1547
+ i2 = i0 + 2
1548
+ i3 = i1 + 2
1549
+ ss = si, sj
1550
+
1551
+ array = var.data[ss]
1552
+ if np.ndim(var.grid.data[0]) == 1:
1553
+ x = var.grid.data[i0][si]
1554
+ y = var.grid.data[i1][sj]
1555
+ else:
1556
+ x = var.grid.data[i0][ss]
1557
+ y = var.grid.data[i1][ss]
1558
+
1559
+ idx = [i0, i1, i2, i3]
1560
+
1561
+ if scale > 0:
1562
+ xscale = scale
1563
+ yscale = scale
1564
+
1565
+ extents = list(var.grid.extents)
1566
+ # This is an attempt to get both axes to use the same scale
1567
+ # I'm not at all confident that it always does the right thing
1568
+ if xscale == 0 and yscale == 0:
1569
+ xlen = max(abs(extents[i2]), abs(extents[i0]))
1570
+ ylen = max(abs(extents[i3]), abs(extents[i1]))
1571
+ ratio1 = ylen / xlen
1572
+ ratio2 = xlen / ylen
1573
+ if ratio1 >= 1 and ratio1 < 1e3:
1574
+ xscale = xlen
1575
+ yscale = xlen
1576
+ elif ratio2 >= 1 and ratio2 < 1e3:
1577
+ xscale = ylen
1578
+ yscale = ylen
1579
+
1580
+ if xscale == 0:
1581
+ length = max(abs(extents[i2]), abs(extents[i0]))
1582
+ mult_x, sym_x = get_si_prefix(length)
1583
+ else:
1584
+ mult_x, sym_x = get_si_prefix(xscale)
1585
+
1586
+ if yscale == 0:
1587
+ length = max(abs(extents[i3]), abs(extents[i1]))
1588
+ mult_y, sym_y = get_si_prefix(length)
1589
+ else:
1590
+ mult_y, sym_y = get_si_prefix(yscale)
1591
+
1592
+ xlabel = var.grid.labels[i0] \
1593
+ + ' $(' + escape_latex(sym_x + var.grid.units[i0]) + ')$'
1594
+ ylabel = var.grid.labels[i1] \
1595
+ + ' $(' + escape_latex(sym_y + var.grid.units[i1]) + ')$'
1596
+
1597
+ var_label = var.name + ' $(' + escape_latex(var.units) + ')$'
1598
+
1599
+ if hasattr(var, 'blocklist'):
1600
+ data = var.blocklist
1601
+
1602
+ plot2d_array(array=array, x=x, y=y, extents=extents, var_label=var_label,
1603
+ xlabel=xlabel, ylabel=ylabel, idx=idx, iso=iso, fast=fast,
1604
+ title=title, full=full, vrange=vrange, reflect=reflect,
1605
+ norm=norm, hold=hold, xscale=xscale, yscale=yscale,
1606
+ scale=scale, figure=figure, subplot=subplot,
1607
+ add_cbar=add_cbar, cbar_label=cbar_label, cbar_top=cbar_top,
1608
+ **kwargs)
1609
+
1610
+
1611
+ def plot2d_update(var):
1612
+ global fig, im
1613
+
1614
+ im.set_array(var.ravel())
1615
+ im.autoscale()
1616
+ fig.canvas.draw()
1617
+
1618
+
1619
+ def plot_levels(var, r0=None, r1=None, nl=10, iso=None, out=False,
1620
+ title=True, levels=True, **kwargs):
1621
+ global data
1622
+
1623
+ try:
1624
+ plt.clf()
1625
+ except:
1626
+ pass
1627
+
1628
+ if iso is None:
1629
+ if hasattr(var, 'blocklist'):
1630
+ data = var.blocklist
1631
+ iso = get_default_iso(data)
1632
+
1633
+ if np.ndim(var.grid.data[0]) == 1:
1634
+ X, Y = np.meshgrid(var.grid.data[1], var.grid.data[0])
1635
+ else:
1636
+ if tuple(var.grid.dims) == tuple(var.dims):
1637
+ X = var.grid.data[0]
1638
+ Y = var.grid.data[1]
1639
+ else:
1640
+ X = var.grid_mid.data[0]
1641
+ Y = var.grid_mid.data[1]
1642
+
1643
+ if r0 is None:
1644
+ r0 = np.min(var.data)
1645
+ r1 = np.max(var.data)
1646
+ dr = (r1 - r0) / (nl + 1)
1647
+ r0 += dr
1648
+ r1 -= dr
1649
+
1650
+ rl = r0 + 1.0 * (r1 - r0) * np.array(range(nl)) / (nl - 1)
1651
+
1652
+ fig = plt.gcf()
1653
+ if out:
1654
+ gs = plt.GridSpec(1, 1) # , width_ratios=[8, 1])
1655
+ ax = plt.subplot(gs[0])
1656
+ else:
1657
+ ax = plt.gca()
1658
+
1659
+ k = 'colors'
1660
+ if k not in kwargs:
1661
+ kwargs[k] = 'k'
1662
+ k = 'linewidths'
1663
+ if k not in kwargs:
1664
+ kwargs[k] = 0.5
1665
+ k = 'levels'
1666
+ if k not in kwargs:
1667
+ kwargs[k] = rl
1668
+
1669
+ cs = ax.contour(X, Y, var.data, **kwargs)
1670
+
1671
+ if levels:
1672
+ fmt = {}
1673
+ for l, i in zip(cs.levels, range(1, len(cs.levels)+1)):
1674
+ fmt[l] = str(i)
1675
+
1676
+ sidx = ""
1677
+ slvl = ""
1678
+ for l, i in reversed(list(zip(cs.levels, range(1, len(cs.levels)+1)))):
1679
+ # sidx += rtn + "%i" % i
1680
+ # slvl += rtn + "%-6.4g" % l
1681
+ # rtn = "\n"
1682
+ sidx += "%i\n" % i
1683
+ slvl += "%-6.4g\n" % l
1684
+
1685
+ t1 = TextArea('Level', textprops=dict(color='k', fontsize='small'))
1686
+ t2 = TextArea(sidx, textprops=dict(color='k', fontsize='small'))
1687
+ tl = VPacker(children=[t1, t2], align="center", pad=0, sep=2)
1688
+
1689
+ lname = var.name.replace("_node", "")
1690
+ t3 = TextArea(lname, textprops=dict(color='k', fontsize='small'))
1691
+ t4 = TextArea(slvl, textprops=dict(color='k', fontsize='small'))
1692
+ tr = VPacker(children=[t3, t4], align="center", pad=0, sep=2)
1693
+
1694
+ t = HPacker(children=[tl, tr], align="center", pad=0, sep=8)
1695
+
1696
+ if out:
1697
+ t = AnchoredOffsetbox(loc=2, child=t, pad=.4,
1698
+ bbox_to_anchor=(1.01, 1), frameon=True,
1699
+ bbox_transform=ax.transAxes, borderpad=0)
1700
+ else:
1701
+ t = AnchoredOffsetbox(loc=1, child=t, pad=.4,
1702
+ bbox_to_anchor=(1, 1), frameon=True,
1703
+ bbox_transform=ax.transAxes, borderpad=.4)
1704
+ t.set_clip_on(False)
1705
+ ax.add_artist(t)
1706
+
1707
+ plt.clabel(cs, cs.levels, fmt=fmt, inline_spacing=2, fontsize=8)
1708
+
1709
+ ax.set_xlabel(var.grid.labels[0] + ' $('
1710
+ + escape_latex(var.grid.units[0]) + ')$')
1711
+ ax.set_ylabel(var.grid.labels[1] + ' $('
1712
+ + escape_latex(var.grid.units[1]) + ')$')
1713
+
1714
+ if title:
1715
+ if out:
1716
+ # suptitle(get_title(), fontsize='large')
1717
+ plt.suptitle(get_title(var), fontsize='large', y=0.92)
1718
+ else:
1719
+ plt.title(get_title(var), fontsize='large', y=1.03)
1720
+
1721
+ plt.axis('tight')
1722
+ if iso:
1723
+ plt.axis('image')
1724
+
1725
+ plt.draw()
1726
+ if out:
1727
+ gs.tight_layout(fig, rect=[0, -0.01, 0.95, 0.92])
1728
+ # fw = fig.get_window_extent().width
1729
+ # tw = fw*.06+t1.get_children()[0].get_window_extent().width \
1730
+ # + t3.get_children()[0].get_window_extent().width
1731
+ # print(1-tw/fw)
1732
+ # gs.update(right=1-tw/fw)
1733
+ # ax.set_position([box.x0, box.y0, box.width + bw, box.height])
1734
+ else:
1735
+ fig.set_tight_layout(True)
1736
+ plt.draw()
1737
+
1738
+
1739
+ def plot_contour(var, r0=None, r1=None, nl=10, iso=None, title=True, **kwargs):
1740
+ return plot_levels(var, r0=r0, r1=r1, nl=nl, iso=iso, out=False,
1741
+ title=title, levels=False, **kwargs)
1742
+
1743
+
1744
+ def getdata(fname, wkd=None, verbose=True, squeeze=False):
1745
+ global data, t, step, p, ppi, ppe, rho, ei, ee, vx, vy, vz, bx, by, bz
1746
+ global x, y, z, xc, yc, zc, grid, grid_mid
1747
+ global old_mtime, old_filename, old_size, cached
1748
+ global wkdir
1749
+
1750
+ if wkd is not None:
1751
+ wkdir = wkd
1752
+
1753
+ if isinstance(fname, int):
1754
+ try:
1755
+ filename = wkdir + "/%0.4i.sdf" % fname
1756
+ except:
1757
+ filename = wkdir + "/" + fname
1758
+ else:
1759
+ filename = fname
1760
+
1761
+ try:
1762
+ st = os.stat(filename)
1763
+ except OSError as e:
1764
+ try:
1765
+ filename = "./%0.4i.sdf" % fname
1766
+ except:
1767
+ filename = "./" + fname
1768
+ try:
1769
+ st = os.stat(filename)
1770
+ wkdir = '.'
1771
+ except OSError as e:
1772
+ print("ERROR opening file {0}: {1}".format(filename, e.strerror))
1773
+ raise
1774
+
1775
+ if st.st_mtime != old_mtime or st.st_size != old_size \
1776
+ or filename != old_filename:
1777
+ if verbose:
1778
+ print("Reading file " + filename)
1779
+ data = sdf.read(filename)
1780
+ old_mtime = st.st_mtime
1781
+ old_size = st.st_size
1782
+ old_filename = filename
1783
+ else:
1784
+ cached = True
1785
+ return data
1786
+
1787
+ cached = False
1788
+
1789
+ if squeeze:
1790
+ for key, value in data.__dict__.items():
1791
+ # Remove single element dimensions
1792
+ try:
1793
+ dims = []
1794
+ for element in value.dims:
1795
+ dims.append([0, element-1])
1796
+ subarray(value, dims)
1797
+ except:
1798
+ pass
1799
+
1800
+ sdfdict = {}
1801
+ for key, value in data.__dict__.items():
1802
+ if hasattr(value, "id"):
1803
+ sdfdict[value.id] = value
1804
+ else:
1805
+ sdfdict[key] = value
1806
+
1807
+ fdict = {}
1808
+ table = {'time': 't'}
1809
+ k = 'Header'
1810
+ if k in sdfdict:
1811
+ h = sdfdict[k]
1812
+ k = list(table.keys())[0]
1813
+ if k in h:
1814
+ key = table[k]
1815
+ var = h[k]
1816
+ if verbose:
1817
+ print(key + str(np.shape(var)) + ' = ' + k)
1818
+ fdict[key] = var
1819
+ globals()[key] = var
1820
+ builtins.__dict__[key] = var
1821
+
1822
+ table = {'Pressure': 'p',
1823
+ 'Pressure_ion': 'ppi',
1824
+ 'Pressure_electron': 'ppe',
1825
+ 'Rho': 'rho',
1826
+ 'Energy_ion': 'ei',
1827
+ 'Energy_electron': 'ee',
1828
+ 'Vx': 'vx',
1829
+ 'Vy': 'vy',
1830
+ 'Vz': 'vz',
1831
+ 'Vr': 'vx',
1832
+ 'VTheta': 'vz',
1833
+ 'Bx': 'bx',
1834
+ 'By': 'by',
1835
+ 'Bz': 'bz',
1836
+ 'Br': 'bx',
1837
+ 'Bt': 'bz',
1838
+ 'bx': 'bx',
1839
+ 'by': 'by',
1840
+ 'bz': 'bz',
1841
+ 'ex': 'ex',
1842
+ 'ey': 'ey',
1843
+ 'ez': 'ez',
1844
+ 'jx': 'jx',
1845
+ 'jy': 'jy',
1846
+ 'jz': 'jz'}
1847
+
1848
+ rz = False
1849
+ if 'Vr' in sdfdict:
1850
+ rz = True
1851
+
1852
+ if rz:
1853
+ table['Vz'] = 'vy'
1854
+ table['Bz'] = 'by'
1855
+
1856
+ inv_table = {}
1857
+ for k, v in table.items():
1858
+ inv_table[v] = inv_table.get(v, [])
1859
+ inv_table[v].append(k)
1860
+
1861
+ for k in table:
1862
+ if k in sdfdict:
1863
+ key = table[k]
1864
+ if hasattr(sdfdict[k], "data"):
1865
+ var = sdfdict[k].data
1866
+ else:
1867
+ var = sdfdict[k]
1868
+ dims = str(tuple(int(i) for i in sdfdict[k].dims))
1869
+ if verbose:
1870
+ print(key + dims + ' = ' + k)
1871
+ fdict[key] = var
1872
+ globals()[key] = var
1873
+ builtins.__dict__[key] = var
1874
+
1875
+ k = 'grid'
1876
+ if k in sdfdict:
1877
+ vargrid = sdfdict[k]
1878
+ grid = vargrid
1879
+ keys = 'x', 'y', 'z'
1880
+ for n in range(np.size(vargrid.dims)):
1881
+ key = keys[n]
1882
+ var = vargrid.data[n]
1883
+ dims = str(tuple(int(i) for i in sdfdict[k].dims))
1884
+ if verbose:
1885
+ print(key + dims + ' = ' + k)
1886
+ fdict[key] = var
1887
+ globals()[key] = var
1888
+ builtins.__dict__[key] = var
1889
+
1890
+ k = 'grid_mid'
1891
+ if k in sdfdict:
1892
+ vargrid = sdfdict[k]
1893
+ grid_mid = vargrid
1894
+ keys = 'xc', 'yc', 'zc'
1895
+ for n in range(np.size(vargrid.dims)):
1896
+ key = keys[n]
1897
+ var = vargrid.data[n]
1898
+ dims = str(tuple(int(i) for i in sdfdict[k].dims))
1899
+ if verbose:
1900
+ print(key + dims + ' = ' + k)
1901
+ fdict[key] = var
1902
+ globals()[key] = var
1903
+ builtins.__dict__[key] = var
1904
+
1905
+ # Export particle arrays
1906
+ for k, value in data.__dict__.items():
1907
+ if type(value) != sdf.BlockPointVariable \
1908
+ and type(value) != sdf.BlockPointMesh:
1909
+ continue
1910
+ key = re.sub(r'[^a-z0-9]', '_', value.id.lower())
1911
+ if hasattr(value, "data"):
1912
+ var = value.data
1913
+ else:
1914
+ var = value
1915
+ dims = str(tuple(int(i) for i in value.dims))
1916
+ if type(value) == sdf.BlockPointVariable:
1917
+ if verbose:
1918
+ print(key + dims + ' = ' + value.name)
1919
+ fdict[key] = var
1920
+ globals()[key] = var
1921
+ builtins.__dict__[key] = var
1922
+ else:
1923
+ vargrid = value
1924
+ grid = vargrid
1925
+ keys = 'x', 'y', 'z'
1926
+ for n in range(np.size(value.dims)):
1927
+ gkey = keys[n] + '_' + key
1928
+ var = value.data[n]
1929
+ dims = str(tuple(int(i) for i in value.dims))
1930
+ if verbose:
1931
+ print(gkey + dims + ' = ' + k + ' ' + keys[n])
1932
+ fdict[gkey] = var
1933
+ globals()[gkey] = var
1934
+ builtins.__dict__[gkey] = var
1935
+
1936
+ # X, Y = np.meshgrid(x, y)
1937
+ return data
1938
+
1939
+
1940
+ def ogrid(skip=None, **kwargs):
1941
+ global x, y, mult_x, mult_y
1942
+ if np.ndim(x) == 1:
1943
+ X, Y = np.meshgrid(x, y)
1944
+ else:
1945
+ s = slice(None, None, skip)
1946
+ X = x[s, s]
1947
+ Y = y[s, s]
1948
+ X = np.multiply(mult_x, X)
1949
+ Y = np.multiply(mult_y, Y)
1950
+ k = 'lw'
1951
+ if k not in kwargs and 'linewidth' not in kwargs:
1952
+ kwargs[k] = 0.5
1953
+ k = 'color'
1954
+ if k not in kwargs:
1955
+ kwargs[k] = 'k'
1956
+ plt.plot(X, Y, **kwargs)
1957
+ plt.plot(X.transpose(), Y.transpose(), **kwargs)
1958
+
1959
+
1960
+ def plotgrid(fname=None, iso=None, title=True):
1961
+ if type(fname) is sdf.BlockList or type(fname) is dict:
1962
+ dat = fname
1963
+ elif fname is not None:
1964
+ dat = getdata(fname, verbose=verbose)
1965
+
1966
+ if iso is None:
1967
+ iso = get_default_iso(dat)
1968
+
1969
+ ogrid()
1970
+
1971
+ ax = plt.gca()
1972
+
1973
+ ax.set_xlabel(grid.labels[0] + ' $(' + escape_latex(grid.units[0]) + ')$')
1974
+ ax.set_ylabel(grid.labels[1] + ' $(' + escape_latex(grid.units[1]) + ')$')
1975
+
1976
+ if title:
1977
+ plt.title(get_title(), fontsize='large', y=1.03)
1978
+
1979
+ plt.axis('tight')
1980
+ if iso:
1981
+ plt.axis('image')
1982
+
1983
+ plt.draw()
1984
+
1985
+ fig = plt.gcf()
1986
+ fig.set_tight_layout(True)
1987
+ plt.draw()
1988
+
1989
+
1990
+ def axis_offset(boxed=False):
1991
+ ax = plt.gca()
1992
+ xlab = ax.get_xlabel()
1993
+ ylab = ax.get_ylabel()
1994
+
1995
+ f = 1e-3
1996
+
1997
+ # for o in ax.findobj():
1998
+ for l in ax.get_lines():
1999
+ bb = l.get_clip_box()
2000
+ bb._bbox = Bbox([[-f, -f], [1+2*f, 1+2*f]])
2001
+ l.set_clip_box(bb)
2002
+ # l.set_clip_on(False)
2003
+
2004
+ if boxed:
2005
+ r = matplotlib.patches.Rectangle((-f, -f), 1+2*f, 1+2*f,
2006
+ transform=ax.transAxes)
2007
+ r.set_color((0, 0, 0, 0))
2008
+ r.set_edgecolor('k')
2009
+ r.set_clip_on(False)
2010
+ ax.add_patch(r)
2011
+
2012
+ w = 1.1
2013
+ gap = 8
2014
+ ax.spines['top'].set_visible(False)
2015
+ ax.spines['right'].set_visible(False)
2016
+ ax.spines['left'].set_position(('outward', gap))
2017
+ ax.spines['left'].set_linewidth(w)
2018
+ ax.spines['bottom'].set_position(('outward', gap))
2019
+ ax.spines['bottom'].set_linewidth(w)
2020
+ ax.tick_params(direction='out', width=w, length=4.5, top='off',
2021
+ right='off')
2022
+ ax.set_xlabel(xlab)
2023
+ ax.set_ylabel(ylab)
2024
+
2025
+ plt.draw()
2026
+
2027
+
2028
+ def tuple_to_slice(slices):
2029
+ subscripts = []
2030
+ for val in slices:
2031
+ start = val[0]
2032
+ end = val[1]
2033
+ if end is not None:
2034
+ end = end + 1
2035
+ subscripts.append(slice(start, end, 1))
2036
+ subscripts = tuple(subscripts)
2037
+ return subscripts
2038
+
2039
+
2040
+ def subarray(base, slices):
2041
+ if (len(slices) != len(base.dims)):
2042
+ print("Must specify a range in all dimensions")
2043
+ return None
2044
+ dims = []
2045
+ # Construct the lengths of the subarray
2046
+ for x in range(0, len(slices)):
2047
+ begin = slices[x][0]
2048
+ end = slices[x][1]
2049
+ if begin is None:
2050
+ begin = 0
2051
+ if end is None:
2052
+ end = base.dims[x]
2053
+ if (end-begin != 0):
2054
+ dims.append(end-begin+1)
2055
+
2056
+ subscripts = tuple_to_slice(slices)
2057
+ base.data = np.squeeze(base.data[subscripts])
2058
+ base.dims = tuple(dims)
2059
+
2060
+
2061
+ def list_variables(data):
2062
+ dct = data.__dict__
2063
+ for key in sorted(dct):
2064
+ try:
2065
+ val = dct[key]
2066
+ print('{} {} {}'.format(key, type(val),
2067
+ np.array2string(np.array(val.dims), separator=', ')))
2068
+ except:
2069
+ pass
2070
+
2071
+
2072
+ def escape_latex(string):
2073
+ return string.replace('%', '\%')
2074
+
2075
+
2076
+ pi = 3.141592653589793238462643383279503
2077
+ q0 = 1.602176565e-19
2078
+ m0 = 9.10938291e-31
2079
+ c = 2.99792458e8
2080
+ kb = 1.3806488e-23
2081
+ mu0 = pi * 4e-7
2082
+ epsilon0 = 1.0 / mu0 / c**2
2083
+ h_planck = 6.62606957e-34
2084
+ h_bar = h_planck / 2.0 / pi