-
Notifications
You must be signed in to change notification settings - Fork 14
/
Copy pathcore.py
11217 lines (9344 loc) · 499 KB
/
core.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
# This program is distributed under the terms of the GNU General Purpose License (GPL).
# Refer to http://www.gnu.org/licenses/gpl.txt
#
# This file is part of eqtools.
#
# eqtools is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# eqtools is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with eqtools. If not, see <http://www.gnu.org/licenses/>.
"""This module provides the core classes for :py:mod:`eqtools`, including the
base :py:class:`Equilibrium` class.
"""
import scipy
import scipy.interpolate
import scipy.integrate
import scipy.constants
import re
import warnings
# Python 2/3 cross compatibility:
import sys
if sys.version_info.major >= 3:
long = int
# Constants to determine how plot labels are formatted:
B_LABEL = '$B$ [T]'
J_LABEL = '$j$ [MA/m$^2$]'
class ModuleWarning(Warning):
"""Warning class to notify the user of unavailable modules.
"""
pass
try:
from . import trispline
_has_trispline = True
except ImportError:
warnings.warn("trispline module could not be loaded -- tricubic spline "
"interpolation will not be available.",
ModuleWarning)
_has_trispline = False
try:
import matplotlib.pyplot as plt
import matplotlib.widgets as mplw
import matplotlib.gridspec as mplgs
import matplotlib.patches as mpatches
import matplotlib.path as mpath
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib.colorbar import ColorbarBase
from matplotlib.colors import Normalize
from .filewriter import gfile
except Exception:
warnings.warn(
"matplotlib modules could not be loaded -- plotting and gfile"
" writing will not be available.",
ModuleWarning
)
class PropertyAccessMixin(object):
"""Mixin to implement access of getter methods through a property-type
interface without the need to apply a decorator to every property.
For any getter `obj.getSomething()`, the call `obj.Something` will do the
same thing.
This is accomplished by overriding :py:meth:`__getattribute__` such that if
an attribute `ATTR` does not exist it then attempts to call `self.getATTR()`.
If `self.getATTR()` does not exist, an :py:class:`AttributeError` will be
raised as usual.
Also overrides :py:meth:`__setattr__` such that it will raise an
:py:class:`AttributeError` when attempting to write an attribute `ATTR` for
which there is already a method `getATTR`.
"""
def __getattribute__(self, name):
"""Get an attribute.
Tries to get attribute as-written. If this fails, tries to call the
method `get<name>` with no arguments. If this fails, raises
:py:class:`AttributeError`. This effectively generates a Python
'property' for each getter method.
Args:
name (String): Name of the attribute to retrieve. If the instance
has an attribute with this name, the attribute is returned. If
the instance does not have an attribute with this name but does
have a method called 'get'+name, this method is called and the
result is returned.
Returns:
The value of the attribute requested.
Raises:
AttributeError: If neither attribute name or method 'get'+name exist.
"""
try:
return super(Equilibrium, self).__getattribute__(name)
except AttributeError:
try:
return super(Equilibrium, self).__getattribute__('get'+name)()
except AttributeError:
raise AttributeError(
"%(class)s object has no attribute '%(n)s' or method 'get%(n)s'"
% {'class': self.__class__.__name__, 'n': name}
)
def __setattr__(self, name, value):
"""Set an attribute.
Raises :py:class:`AttributeError` if the object already has a method
'get'+name, as creation of such an attribute would interfere with the
automatic property generation in :py:meth:`__getattribute__`.
Args:
name (String): Name of the attribute to set.
value (Object): Value to set the attribute to.
Raises:
AttributeError: If a method called 'get'+name already exists.
"""
if hasattr(self, 'get'+name):
raise AttributeError(
"%(class)s object already has getter method 'get%(n)s', creating "
"attribute '%(n)s' will conflict with automatic property "
"generation." % {'class': self.__class__.__name__, 'n': name}
)
else:
super(Equilibrium, self).__setattr__(name, value)
"""The following is a dictionary to implement length unit conversions. The first
key is the unit are converting FROM, the second the unit you are converting TO.
Supports: m, cm, mm, in, ft, yd, smoot, cubit, hand
"""
_length_conversion = {'m': {'m': 1.0,
'cm': 100.0,
'mm': 1000.0,
'in': 39.37,
'ft': 39.37 / 12.0,
'yd': 39.37 / (12.0 * 3.0),
'smoot': 39.37 / 67.0,
'cubit': 39.37 / 18.0,
'hand': 39.37 / 4.0},
'cm': {'m': 0.01,
'cm': 1.0,
'mm': 10.0,
'in': 39.37 / 100.0,
'ft': 39.37 / (100.0 * 12.0),
'yd': 39.37 / (100.0 * 12.0 * 3.0),
'smoot': 39.37 / (100.0 * 67.0),
'cubit': 39.37 / (100.0 * 18.0),
'hand': 39.37 / (100.0 * 4.0)},
'mm': {'m': 0.001,
'cm': 0.1,
'mm': 1.0,
'in': 39.37 / 1000.0,
'ft': 39.37 / (1000.0 * 12.0),
'yd': 39.37 / (1000.0 * 12.0 * 3.0),
'smoot': 39.37 / (1000.0 * 67.0),
'cubit': 39.37 / (1000.0 * 18.0),
'hand': 39.37 / (1000.0 * 4.0)},
'in': {'m': 1.0 / 39.37,
'cm': 100.0 / 39.37,
'mm': 1000.0 / 39.37,
'in': 1.0,
'ft': 1.0 / 12.0,
'yd': 1.0 / (12.0 * 3.0),
'smoot': 1.0 / 67.0,
'cubit': 1.0 / 18.0,
'hand': 1.0 / 4.0},
'ft': {'m': 12.0 / 39.37,
'cm': 12.0 * 100.0 / 39.37,
'mm': 12.0 * 1000.0 / 39.37,
'in': 12.0,
'ft': 1.0,
'yd': 1.0 / 3.0,
'smoot': 12.0 / 67.0,
'cubit': 12.0 / 18.0,
'hand': 12.0 / 4.0},
'yd': {'m': 3.0 * 12.0 / 39.37,
'cm': 3.0 * 12.0 * 100.0 / 39.37,
'mm': 3.0 * 12.0 * 1000.0 / 39.37,
'in': 3.0 * 12.0,
'ft': 3.0,
'yd': 1.0,
'smoot': 3.0 * 12.0 / 67.0,
'cubit': 3.0 * 12.0 / 18.0,
'hand': 3.0 * 12.0 / 4.0},
'smoot': {'m': 67.0 / 39.37,
'cm': 67.0 * 100.0 / 39.37,
'mm': 67.0 * 1000.0 / 39.37,
'in': 67.0,
'ft': 67.0 / 12.0,
'yd': 67.0 / (12.0 * 3.0),
'smoot': 1.0,
'cubit': 67.0 / 18.0,
'hand': 67.0 / 4.0},
'cubit': {'m': 18.0 / 39.37,
'cm': 18.0 * 100.0 / 39.37,
'mm': 18.0 * 1000.0 / 39.37,
'in': 18.0,
'ft': 18.0 / 12.0,
'yd': 18.0 / (12.0 * 3.0),
'smoot': 18.0 / 67.0,
'cubit': 1.0,
'hand': 18.0 / 4.0},
'hand': {'m': 4.0 / 39.37,
'cm': 4.0 * 100.0 / 39.37,
'mm': 4.0 * 1000.0 / 39.37,
'in': 4.0,
'ft': 4.0 / 12.0,
'yd': 4.0 / (12.0 * 3.0),
'smoot': 4.0 / 67.0,
'cubit': 4.0 / 18.0,
'hand': 1.0}}
def inPolygon(polyx, polyy, pointx, pointy):
"""Function calculating whether a given point is within a 2D polygon.
Given an array of X,Y coordinates describing a 2D polygon, checks whether a
point given by x,y coordinates lies within the polygon. Operates via a
ray-casting approach - the function projects a semi-infinite ray parallel to
the positive horizontal axis, and counts how many edges of the polygon this
ray intersects. For a simply-connected polygon, this determines whether the
point is inside (even number of crossings) or outside (odd number of
crossings) the polygon, by the Jordan Curve Theorem.
Args:
polyx (Array-like): Array of x-coordinates of the vertices of the polygon.
polyy (Array-like): Array of y-coordinates of the vertices of the polygon.
pointx (Int or float): x-coordinate of test point.
pointy (Int or float): y-coordinate of test point.
Returns:
result (Boolean): True/False result for whether the point is contained within the polygon.
"""
# generator function for "lines" - pairs of (x,y) coords describing each edge of the polygon.
def lines():
p0x = polyx[-1]
p0y = polyy[-1]
p0 = (p0x, p0y)
for i, x in enumerate(polyx):
y = polyy[i]
p1 = (x, y)
yield p0, p1
p0 = p1
result = False
for p0, p1 in lines():
if (
(p0[1] > pointy) != (p1[1] > pointy)
) and (
pointx < ((p1[0]-p0[0])*(pointy-p0[1])/(p1[1]-p0[1]) + p0[0])
):
result = not result
return result
class Equilibrium(object):
"""Abstract class of data handling object for magnetic reconstruction outputs.
Defines the mapping routines and method fingerprints necessary. Each
variable or set of variables is recovered with a corresponding getter method.
Essential data for mapping are pulled on initialization (psirz grid, for
example) to frontload overhead. Additional data are pulled at the first
request and stored for subsequent usage.
.. note:: This abstract class should not be used directly. Device- and code-
specific subclasses are set up to account for inter-device/-code
differences in data storage.
Keyword Args:
length_unit (String): Sets the base unit used for any quantity whose
dimensions are length to any power. Valid options are:
=========== ===========================================================================================
'm' meters
'cm' centimeters
'mm' millimeters
'in' inches
'ft' feet
'yd' yards
'smoot' smoots
'cubit' cubits
'hand' hands
'default' whatever the default in the tree is (no conversion is performed, units may be inconsistent)
=========== ===========================================================================================
Default is 'm' (all units taken and returned in meters).
tspline (Boolean): Sets whether or not interpolation in time is
performed using a tricubic spline or nearest-neighbor interpolation.
Tricubic spline interpolation requires at least four complete
equilibria at different times. It is also assumed that they are
functionally correlated, and that parameters do not vary out of
their boundaries (derivative = 0 boundary condition). Default is
False (use nearest-neighbor interpolation).
monotonic (Boolean): Sets whether or not the "monotonic" form of time
window finding is used. If True, the timebase must be monotonically
increasing. Default is False (use slower, safer method).
verbose (Boolean): Allows or blocks console readout during operation.
Defaults to True, displaying useful information for the user. Set to
False for quiet usage or to avoid console clutter for multiple
instances.
Raises:
ValueError: If `length_unit` is not a valid unit specifier.
ValueError: If `tspline` is True but module trispline did not load
successfully.
"""
def __init__(
self, length_unit='m', tspline=False, monotonic=True, verbose=True
):
if (
(length_unit != 'default') and
(not (length_unit in _length_conversion))
):
raise ValueError(
"Unit '%s' not a valid unit specifier!" % length_unit
)
else:
self._length_unit = length_unit
self._tricubic = bool(tspline)
self._monotonic = bool(monotonic)
self._verbose = bool(verbose)
# These are indexes of splines, and become higher dimensional splines
# with the setting of the tspline keyword.
self._psiOfRZSpline = {}
self._phiNormSpline = {}
self._volNormSpline = {}
self._RmidSpline = {}
self._magRSpline = {}
self._magZSpline = {}
self._RmidOutSpline = {}
self._psiOfPsi0Spline = {}
self._psiOfLCFSSpline = {}
self._RmidToPsiNormSpline = {}
self._phiNormToPsiNormSpline = {}
self._volNormToPsiNormSpline = {}
self._AOutSpline = {}
self._qSpline = {}
self._FSpline = {}
self._FToPsinormSpline = {}
self._FFPrimeSpline = {}
self._pSpline = {}
self._pPrimeSpline = {}
self._vSpline = {}
self._BtVacSpline = {}
def __str__(self):
"""String representation of this instance.
Returns:
string (String): String describing this object.
"""
return 'This is an abstract class. Please use machine-specific subclass.'
def __getstate__(self):
"""Deletes all of the stored splines, since they aren't pickleable.
"""
self._psiOfRZSpline = {}
self._phiNormSpline = {}
self._volNormSpline = {}
self._RmidSpline = {}
self._magRSpline = {}
self._magZSpline = {}
self._RmidOutSpline = {}
self._psiOfPsi0Spline = {}
self._psiOfLCFSSpline = {}
self._RmidToPsiNormSpline = {}
self._phiNormToPsiNormSpline = {}
self._volNormToPsiNormSpline = {}
self._AOutSpline = {}
self._qSpline = {}
self._FSpline = {}
self._FToPsinormSpline = {}
self._FFPrimeSpline = {}
self._pSpline = {}
self._pPrimeSpline = {}
self._vSpline = {}
self._BtVacSpline = {}
return self.__dict__
####################
# Mapping routines #
####################
def rho2rho(self, origin, destination, *args, **kwargs):
r"""Convert from one coordinate to another.
Args:
origin (String): Indicates which coordinates the data are given in.
Valid options are:
======= ========================
RZ R,Z coordinates
psinorm Normalized poloidal flux
phinorm Normalized toroidal flux
volnorm Normalized volume
Rmid Midplane major radius
r/a Normalized minor radius
======= ========================
Additionally, each valid option may be prepended with 'sqrt'
to specify the square root of the desired unit.
destination (String): Indicates which coordinates to convert to.
Valid options are:
======= =================================
psinorm Normalized poloidal flux
phinorm Normalized toroidal flux
volnorm Normalized volume
Rmid Midplane major radius
r/a Normalized minor radius
q Safety factor
F Flux function :math:`F=RB_{\phi}`
FFPrime Flux function :math:`FF'`
p Pressure
pprime Pressure gradient
v Flux surface volume
======= =================================
Additionally, each valid option may be prepended with 'sqrt'
to specify the square root of the desired unit.
rho (Array-like or scalar float): Values of the starting coordinate
to map to the new coordinate. Will be two arguments `R`, `Z` if
`origin` is 'RZ'.
t (Array-like or scalar float): Times to perform the conversion at.
If `t` is a single value, it is used for all of the elements of
`rho`. If the `each_t` keyword is True, then `t` must be scalar
or have exactly one dimension. If the `each_t` keyword is False,
`t` must have the same shape as `rho` (or the meshgrid of `R`
and `Z` if `make_grid` is True).
Keyword Args:
sqrt (Boolean): Set to True to return the square root of `rho`. Only
the square root of positive values is taken. Negative values are
replaced with zeros, consistent with Steve Wolfe's IDL
implementation efit_rz2rho.pro. Default is False.
each_t (Boolean): When True, the elements in `rho` are evaluated at
each value in `t`. If True, `t` must have only one dimension (or
be a scalar). If False, `t` must match the shape of `rho` or be
a scalar. Default is True (evaluate ALL `rho` at EACH element in
`t`).
make_grid (Boolean): Only applicable if `origin` is 'RZ'. Set to
True to pass `R` and `Z` through :py:func:`scipy.meshgrid`
before evaluating. If this is set to True, `R` and `Z` must each
only have a single dimension, but can have different lengths.
Default is False (do not form meshgrid).
rho (Boolean): Set to True to return r/a (normalized minor radius)
instead of Rmid when `destination` is Rmid. Default is False
(return major radius, Rmid).
length_unit (String or 1): Length unit that quantities are
given/returned in, as applicable. If a string is given, it must
be a valid unit specifier:
=========== ===========
'm' meters
'cm' centimeters
'mm' millimeters
'in' inches
'ft' feet
'yd' yards
'smoot' smoots
'cubit' cubits
'hand' hands
'default' meters
=========== ===========
If length_unit is 1 or None, meters are assumed. The default
value is 1 (use meters).
k (positive int): The degree of polynomial spline interpolation to
use in converting coordinates.
return_t (Boolean): Set to True to return a tuple of (`rho`,
`time_idxs`), where `time_idxs` is the array of time indices
actually used in evaluating `rho` with nearest-neighbor
interpolation. (This is mostly present as an internal helper.)
Default is False (only return `rho`).
Returns:
`rho` or (`rho`, `time_idxs`)
* **rho** (`Array or scalar float`) - The converted coordinates. If
all of the input arguments are scalar, then a scalar is returned.
Otherwise, a scipy Array is returned.
* **time_idxs** (Array with same shape as `rho`) - The indices
(in :py:meth:`self.getTimeBase`) that were used for
nearest-neighbor interpolation. Only returned if `return_t` is
True.
Raises:
ValueError: If `origin` is not one of the supported values.
Examples:
All assume that `Eq_instance` is a valid instance of the appropriate
extension of the :py:class:`Equilibrium` abstract class.
Find single psinorm value at r/a=0.6, t=0.26s::
psi_val = Eq_instance.rho2rho('r/a', 'psinorm', 0.6, 0.26)
Find psinorm values at r/a points 0.6 and 0.8 at the
single time t=0.26s::
psi_arr = Eq_instance.rho2rho('r/a', 'psinorm', [0.6, 0.8], 0.26)
Find psinorm values at r/a of 0.6 at times t=[0.2s, 0.3s]::
psi_arr = Eq_instance.rho2rho('r/a', 'psinorm', 0.6, [0.2, 0.3])
Find psinorm values at (r/a, t) points (0.6, 0.2s) and (0.5, 0.3s)::
psi_arr = Eq_instance.rho2rho('r/a', 'psinorm', [0.6, 0.5], [0.2, 0.3], each_t=False)
"""
if origin.startswith('sqrt'):
args = list(args)
args[0] = scipy.asarray(args[0])**2
origin = origin[4:]
if destination.startswith('sqrt'):
kwargs['sqrt'] = True
destination = destination[4:]
if origin == 'RZ':
return self.rz2rho(destination, *args, **kwargs)
elif origin == 'Rmid':
return self.rmid2rho(destination, *args, **kwargs)
elif origin == 'r/a':
return self.roa2rho(destination, *args, **kwargs)
elif origin == 'psinorm':
return self.psinorm2rho(destination, *args, **kwargs)
elif origin == 'phinorm':
return self.phinorm2rho(destination, *args, **kwargs)
elif origin == 'volnorm':
return self.volnorm2rho(destination, *args, **kwargs)
else:
raise ValueError(
"rho2rho: Unsupported origin coordinate method '%s'!" % origin
)
def rz2psi(
self, R, Z, t, return_t=False, make_grid=False, each_t=True,
length_unit=1
):
r"""Converts the passed R, Z, t arrays to psi (unnormalized poloidal flux) values.
What is usually returned by EFIT is the stream function,
:math:`\psi=\psi_p/(2\pi)` which has units of Wb/rad.
Args:
R (Array-like or scalar float): Values of the radial coordinate to
map to poloidal flux. If `R` and `Z` are both scalar values,
they are used as the coordinate pair for all of the values in
`t`. Must have the same shape as `Z` unless the `make_grid`
keyword is set. If the `make_grid` keyword is True, `R` must
have exactly one dimension.
Z (Array-like or scalar float): Values of the vertical coordinate to
map to poloidal flux. If `R` and `Z` are both scalar values,
they are used as the coordinate pair for all of the values in
`t`. Must have the same shape as `R` unless the `make_grid`
keyword is set. If the `make_grid` keyword is True, `Z` must
have exactly one dimension.
t (Array-like or scalar float): Times to perform the conversion at.
If `t` is a single value, it is used for all of the elements of
`R`, `Z`. If the `each_t` keyword is True, then `t` must be
scalar or have exactly one dimension. If the `each_t` keyword is
False, `t` must have the same shape as `R` and `Z` (or their
meshgrid if `make_grid` is True).
Keyword Args:
each_t (Boolean): When True, the elements in `R`, `Z` are evaluated
at each value in `t`. If True, `t` must have only one dimension
(or be a scalar). If False, `t` must match the shape of `R` and
`Z` or be a scalar. Default is True (evaluate ALL `R`, `Z` at
EACH element in `t`).
make_grid (Boolean): Set to True to pass `R` and `Z` through
:py:func:`scipy.meshgrid` before evaluating. If this is set to
True, `R` and `Z` must each only have a single dimension, but
can have different lengths. Default is False (do not form
meshgrid).
length_unit (String or 1): Length unit that `R`, `Z` are given in.
If a string is given, it must be a valid unit specifier:
=========== ===========
'm' meters
'cm' centimeters
'mm' millimeters
'in' inches
'ft' feet
'yd' yards
'smoot' smoots
'cubit' cubits
'hand' hands
'default' meters
=========== ===========
If length_unit is 1 or None, meters are assumed. The default
value is 1 (use meters).
return_t (Boolean): Set to True to return a tuple of (`rho`,
`time_idxs`), where `time_idxs` is the array of time indices
actually used in evaluating `rho` with nearest-neighbor
interpolation. (This is mostly present as an internal helper.)
Default is False (only return `rho`).
Returns:
`psi` or (`psi`, `time_idxs`)
* **psi** (`Array or scalar float`) - The unnormalized poloidal
flux. If all of the input arguments are scalar, then a scalar is
returned. Otherwise, a scipy Array is returned. If `R` and `Z`
both have the same shape then `psi` has this shape as well,
unless the `make_grid` keyword was True, in which case `psi` has
shape (len(`Z`), len(`R`)).
* **time_idxs** (Array with same shape as `psi`) - The indices
(in :py:meth:`self.getTimeBase`) that were used for
nearest-neighbor interpolation. Only returned if `return_t` is
True.
Examples:
All assume that `Eq_instance` is a valid instance of the appropriate
extension of the :py:class:`Equilibrium` abstract class.
Find single psi value at R=0.6m, Z=0.0m, t=0.26s::
psi_val = Eq_instance.rz2psi(0.6, 0, 0.26)
Find psi values at (R, Z) points (0.6m, 0m) and (0.8m, 0m) at the
single time t=0.26s. Note that the `Z` vector must be fully
specified, even if the values are all the same::
psi_arr = Eq_instance.rz2psi([0.6, 0.8], [0, 0], 0.26)
Find psi values at (R, Z) points (0.6m, 0m) at times t=[0.2s, 0.3s]::
psi_arr = Eq_instance.rz2psi(0.6, 0, [0.2, 0.3])
Find psi values at (R, Z, t) points (0.6m, 0m, 0.2s) and
(0.5m, 0.2m, 0.3s)::
psi_arr = Eq_instance.rz2psi([0.6, 0.5], [0, 0.2], [0.2, 0.3], each_t=False)
Find psi values on grid defined by 1D vector of radial positions `R`
and 1D vector of vertical positions `Z` at time t=0.2s::
psi_mat = Eq_instance.rz2psi(R, Z, 0.2, make_grid=True)
"""
# Check inputs and process into flat arrays with units of meters:
(
R, Z, t, time_idxs, unique_idxs, single_time, single_val,
original_shape
) = self._processRZt(
R,
Z,
t,
make_grid=make_grid,
each_t=each_t,
length_unit=length_unit,
compute_unique=True
)
if self._tricubic:
out_vals = scipy.reshape(
self._getFluxTriSpline().ev(t, Z, R),
original_shape
)
else:
if single_time:
out_vals = self._getFluxBiSpline(time_idxs[0]).ev(Z, R)
if single_val:
out_vals = out_vals[0]
else:
out_vals = scipy.reshape(out_vals, original_shape)
elif each_t:
out_vals = scipy.zeros(
scipy.concatenate(
([len(time_idxs), ], original_shape)
).astype(int)
)
for idx, t_idx in enumerate(time_idxs):
out_vals[idx] = self._getFluxBiSpline(
t_idx
).ev(Z, R).reshape(original_shape)
else:
out_vals = scipy.zeros_like(t, dtype=float)
for t_idx in unique_idxs:
t_mask = (time_idxs == t_idx)
out_vals[t_mask] = self._getFluxBiSpline(
t_idx
).ev(Z[t_mask], R[t_mask])
out_vals = scipy.reshape(out_vals, original_shape)
# Correct for current sign:
out_vals = -1.0 * out_vals * self.getCurrentSign()
if return_t:
if self._tricubic:
return out_vals, (t, single_time, single_val, original_shape)
else:
return out_vals, (
time_idxs, unique_idxs, single_time, single_val,
original_shape
)
else:
return out_vals
def rz2psinorm(self, R, Z, t, return_t=False, sqrt=False, make_grid=False,
each_t=True, length_unit=1):
r"""Calculates the normalized poloidal flux at the given (R, Z, t).
Uses the definition:
.. math::
\texttt{psi\_norm} = \frac{\psi - \psi(0)}{\psi(a) - \psi(0)}
Args:
R (Array-like or scalar float): Values of the radial coordinate to
map to psinorm. If `R` and `Z` are both scalar values,
they are used as the coordinate pair for all of the values in
`t`. Must have the same shape as `Z` unless the `make_grid`
keyword is set. If the `make_grid` keyword is True, `R` must
have exactly one dimension.
Z (Array-like or scalar float): Values of the vertical coordinate to
map to psinorm. If `R` and `Z` are both scalar values,
they are used as the coordinate pair for all of the values in
`t`. Must have the same shape as `R` unless the `make_grid`
keyword is set. If the `make_grid` keyword is True, `Z` must
have exactly one dimension.
t (Array-like or scalar float): Times to perform the conversion at.
If `t` is a single value, it is used for all of the elements of
`R`, `Z`. If the `each_t` keyword is True, then `t` must be
scalar or have exactly one dimension. If the `each_t` keyword is
False, `t` must have the same shape as `R` and `Z` (or their
meshgrid if `make_grid` is True).
Keyword Args:
sqrt (Boolean): Set to True to return the square root of psinorm.
Only the square root of positive values is taken. Negative
values are replaced with zeros, consistent with Steve Wolfe's
IDL implementation efit_rz2rho.pro. Default is False.
each_t (Boolean): When True, the elements in `R`, `Z` are evaluated
at each value in `t`. If True, `t` must have only one dimension
(or be a scalar). If False, `t` must match the shape of `R` and
`Z` or be a scalar. Default is True (evaluate ALL `R`, `Z` at
EACH element in `t`).
make_grid (Boolean): Set to True to pass `R` and `Z` through
:py:func:`scipy.meshgrid` before evaluating. If this is set to
True, `R` and `Z` must each only have a single dimension, but
can have different lengths. Default is False (do not form
meshgrid).
length_unit (String or 1): Length unit that `R`, `Z` are given in.
If a string is given, it must be a valid unit specifier:
=========== ===========
'm' meters
'cm' centimeters
'mm' millimeters
'in' inches
'ft' feet
'yd' yards
'smoot' smoots
'cubit' cubits
'hand' hands
'default' meters
=========== ===========
If length_unit is 1 or None, meters are assumed. The default
value is 1 (use meters).
return_t (Boolean): Set to True to return a tuple of (`rho`,
`time_idxs`), where `time_idxs` is the array of time indices
actually used in evaluating `rho` with nearest-neighbor
interpolation. (This is mostly present as an internal helper.)
Default is False (only return `rho`).
Returns:
`psinorm` or (`psinorm`, `time_idxs`)
* **psinorm** (`Array or scalar float`) - The normalized poloidal
flux. If all of the input arguments are scalar, then a scalar is
returned. Otherwise, a scipy Array is returned. If `R` and `Z`
both have the same shape then `psinorm` has this shape as well,
unless the `make_grid` keyword was True, in which case `psinorm`
has shape (len(`Z`), len(`R`)).
* **time_idxs** (Array with same shape as `psinorm`) - The indices
(in :py:meth:`self.getTimeBase`) that were used for
nearest-neighbor interpolation. Only returned if `return_t` is
True.
Examples:
All assume that `Eq_instance` is a valid instance of the
appropriate extension of the :py:class:`Equilibrium` abstract class.
Find single psinorm value at R=0.6m, Z=0.0m, t=0.26s::
psi_val = Eq_instance.rz2psinorm(0.6, 0, 0.26)
Find psinorm values at (R, Z) points (0.6m, 0m) and (0.8m, 0m) at the
single time t=0.26s. Note that the `Z` vector must be fully specified,
even if the values are all the same::
psi_arr = Eq_instance.rz2psinorm([0.6, 0.8], [0, 0], 0.26)
Find psinorm values at (R, Z) points (0.6m, 0m) at times t=[0.2s, 0.3s]::
psi_arr = Eq_instance.rz2psinorm(0.6, 0, [0.2, 0.3])
Find psinorm values at (R, Z, t) points (0.6m, 0m, 0.2s) and (0.5m, 0.2m, 0.3s)::
psi_arr = Eq_instance.rz2psinorm([0.6, 0.5], [0, 0.2], [0.2, 0.3], each_t=False)
Find psinorm values on grid defined by 1D vector of radial positions `R`
and 1D vector of vertical positions `Z` at time t=0.2s::
psi_mat = Eq_instance.rz2psinorm(R, Z, 0.2, make_grid=True)
"""
psi, blob = self.rz2psi(
R,
Z,
t,
return_t=True,
make_grid=make_grid,
each_t=each_t,
length_unit=length_unit
)
if self._tricubic:
psi_boundary = self._getLCFSPsiSpline()(blob[0]).reshape(blob[-1])
psi_0 = self._getPsi0Spline()(blob[0]).reshape(blob[-1])
else:
psi_boundary = self.getFluxLCFS()[blob[0]]
psi_0 = self.getFluxAxis()[blob[0]]
# If there is more than one time point, we need to expand these
# arrays to be broadcastable:
if not blob[-3]:
if each_t:
for k in range(0, len(blob[-1])):
psi_boundary = scipy.expand_dims(psi_boundary, -1)
psi_0 = scipy.expand_dims(psi_0, -1)
else:
psi_boundary = psi_boundary.reshape(blob[-1])
psi_0 = psi_0.reshape(blob[-1])
psi_norm = (psi - psi_0) / (psi_boundary - psi_0)
if sqrt:
if psi_norm.ndim == 0:
if psi_norm < 0.0:
psi_norm = 0.0
else:
scipy.place(psi_norm, psi_norm < 0, 0)
out = scipy.sqrt(psi_norm)
else:
out = psi_norm
# Unwrap single values to ensure least surprise:
if blob[-2] and blob[-3] and not self._tricubic:
out = out[0]
if return_t:
return out, blob
else:
return out
def rz2phinorm(self, *args, **kwargs):
r"""Calculates the normalized toroidal flux.
Uses the definitions:
.. math::
\texttt{phi} &= \int q(\psi)\,d\psi\\
\texttt{phi\_norm} &= \frac{\phi}{\phi(a)}
This is based on the IDL version efit_rz2rho.pro by Steve Wolfe.
Args:
R (Array-like or scalar float): Values of the radial coordinate to
map to phinorm. If `R` and `Z` are both scalar values,
they are used as the coordinate pair for all of the values in
`t`. Must have the same shape as `Z` unless the `make_grid`
keyword is set. If the `make_grid` keyword is True, `R` must
have exactly one dimension.
Z (Array-like or scalar float): Values of the vertical coordinate to
map to phinorm. If `R` and `Z` are both scalar values,
they are used as the coordinate pair for all of the values in
`t`. Must have the same shape as `R` unless the `make_grid`
keyword is set. If the `make_grid` keyword is True, `Z` must
have exactly one dimension.
t (Array-like or scalar float): Times to perform the conversion at.
If `t` is a single value, it is used for all of the elements of
`R`, `Z`. If the `each_t` keyword is True, then `t` must be
scalar or have exactly one dimension. If the `each_t` keyword is
False, `t` must have the same shape as `R` and `Z` (or their
meshgrid if `make_grid` is True).
Keyword Args:
sqrt (Boolean): Set to True to return the square root of phinorm.
Only the square root of positive values is taken. Negative
values are replaced with zeros, consistent with Steve Wolfe's
IDL implementation efit_rz2rho.pro. Default is False.
each_t (Boolean): When True, the elements in `R`, `Z` are evaluated
at each value in `t`. If True, `t` must have only one dimension
(or be a scalar). If False, `t` must match the shape of `R` and
`Z` or be a scalar. Default is True (evaluate ALL `R`, `Z` at
EACH element in `t`).
make_grid (Boolean): Set to True to pass `R` and `Z` through
:py:func:`scipy.meshgrid` before evaluating. If this is set to
True, `R` and `Z` must each only have a single dimension, but
can have different lengths. Default is False (do not form
meshgrid).
length_unit (String or 1): Length unit that `R`, `Z` are given in.
If a string is given, it must be a valid unit specifier:
=========== ===========
'm' meters
'cm' centimeters
'mm' millimeters
'in' inches
'ft' feet
'yd' yards
'smoot' smoots
'cubit' cubits
'hand' hands
'default' meters
=========== ===========
If length_unit is 1 or None, meters are assumed. The default
value is 1 (use meters).
k (positive int): The degree of polynomial spline interpolation to
use in converting psinorm to phinorm.
return_t (Boolean): Set to True to return a tuple of (`rho`,
`time_idxs`), where `time_idxs` is the array of time indices
actually used in evaluating `rho` with nearest-neighbor
interpolation. (This is mostly present as an internal helper.)
Default is False (only return `rho`).
Returns:
`phinorm` or (`phinorm`, `time_idxs`)
* **phinorm** (`Array or scalar float`) - The normalized toroidal
flux. If all of the input arguments are scalar, then a scalar is
returned. Otherwise, a scipy Array is returned. If `R` and `Z`
both have the same shape then `phinorm` has this shape as well,
unless the `make_grid` keyword was True, in which case `phinorm`
has shape (len(`Z`), len(`R`)).
* **time_idxs** (Array with same shape as `phinorm`) - The indices
(in :py:meth:`self.getTimeBase`) that were used for
nearest-neighbor interpolation. Only returned if `return_t` is
True.
Examples:
All assume that `Eq_instance` is a valid instance of the appropriate
extension of the :py:class:`Equilibrium` abstract class.
Find single phinorm value at R=0.6m, Z=0.0m, t=0.26s::
phi_val = Eq_instance.rz2phinorm(0.6, 0, 0.26)
Find phinorm values at (R, Z) points (0.6m, 0m) and (0.8m, 0m) at the
single time t=0.26s. Note that the `Z` vector must be fully specified,
even if the values are all the same::
phi_arr = Eq_instance.rz2phinorm([0.6, 0.8], [0, 0], 0.26)
Find phinorm values at (R, Z) points (0.6m, 0m) at times t=[0.2s, 0.3s]::
phi_arr = Eq_instance.rz2phinorm(0.6, 0, [0.2, 0.3])
Find phinorm values at (R, Z, t) points (0.6m, 0m, 0.2s) and (0.5m, 0.2m, 0.3s)::
phi_arr = Eq_instance.rz2phinorm([0.6, 0.5], [0, 0.2], [0.2, 0.3], each_t=False)
Find phinorm values on grid defined by 1D vector of radial positions `R`
and 1D vector of vertical positions `Z` at time t=0.2s::
phi_mat = Eq_instance.rz2phinorm(R, Z, 0.2, make_grid=True)
"""
return self._RZ2Quan(self._getPhiNormSpline, *args, **kwargs)
def rz2volnorm(self, *args, **kwargs):
"""Calculates the normalized flux surface volume.