-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathdependencies.jl
3269 lines (2881 loc) · 180 KB
/
dependencies.jl
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
######################################### dependencies.jl ##################################################
# This script defines a large collection of functions that is used by various scripts and apps of the OWCF.
#
# Prior to loading this collection of functions (for example via 'include(extra/dependencies.jl)' when standing
# in the OWCF-folder), you should have performed the usual activation of the OWCF Julia environment.
# This is done as follows
#
# folderpath_OWCF = "/path/to/the/OWCF/folder/"
# using Distributed
# @everywhere begin
# using Pkg
# cd(folderpath_OWCF)
# Pkg.activate(".")
# end
#
# This is performed in e.g. every OWCF start template file. It is also possible to skip the Distributed.jl
# package. This is then done as
#
# folderpath_OWCF = "/path/to/the/OWCF/folder/"
# using Pkg
# cd(folderpath_OWCF)#
# Pkg.activate(".")
#
# Finally, please note that some functions in dependencies.jl might be under construction. Hence, the bad code.
#
# Written by H. Järleblad. Last updated 2024-08-22.
###################################################################################################
println("Loading the Julia packages for the OWCF dependencies... ")
using Distributed
using EFIT
using Equilibrium
using GuidingCenterOrbits
using ForwardDiff
using LinearAlgebra
using OrbitTomography
using Optim
using JLD2
using HDF5
using NetCDF
using FileIO
using Interpolations
using NearestNeighbors
using Statistics
using SharedArrays
using ProgressMeter
using Distributions
using VoronoiDelaunay
import OrbitTomography: orbit_grid
using Base.Iterators
###### Constants needed for dependencies ######
const kB = 1.380649e-23 # Boltzmann constant, J/K
const ϵ0 = 8.8541878128e-12 # Permittivity of free space
###### Structures needed for dependencies ######
"""
grid(r2d,z2d,r,z,phi,nr,nz,nphi)
A structure to represent an (R,z) grid in 2D, together with a
grid in the toroidal phi direction. The r2d and z2d fields are mesh
grids of the (R,z) grid points.
"""
struct grid
r2d::Matrix{Float64}
z2d::Matrix{Float64}
r::Vector{Float64}
z::Vector{Float64}
phi::Vector{Float64}
nr::Int64
nz::Int64
nphi::Int64
end
###### Functions ######
###### Mathematics
"""
gaussian(μ, σ)
gaussian(-||-; mx=μ .+6 .*σ, mn=μ .-6 .*σ, n=50)
Compute the (multi-variate) Gaussian distribution with mean 'μ' and standard deviation 'σ'.
By default, the upper bounds 'mx' of the grid are found by adding 6σ to the mean μ.
The lower bounds 'mn' of the grid are found by subtracting 6σ from the mean μ. By default, the number
of grid points in all dimensions is the same: 50. To use the function to create a D-dimensional Gaussian
distribution, where D is any integer >0, do e.g. the following. The example below is for a 3-dimensional case:
m = [100.0, 0.6, 3.3]
v = [25.0, 0.01, 0.04]
myGauss = gaussian(μ, σ; mx=[200.0, 1.0, 3.8], mn=[0.0, -1.0, 3.0], n=[10,101,104])
The variable 'myGauss' will then be a 10x101x104 array. In the first dimension, the lower and upper bounds will
be 0.0 and 200.0, respectively. And so on for the other dimensions.
The floor_level keyword argument can be used to set all values smaller than floor_level*maximum(f)
to 0.0 before returning the output. f is the Gaussian distribution.
The verbose keyword argument will make the function talk more!
"""
function gaussian(μ::AbstractVector, σ::AbstractVector; mn::AbstractVector=μ .-6 .*σ, mx::AbstractVector=μ .+6 .*σ, n::Union{Int64,Vector{Int64}}=50, floor_level::Float64=0.0, verbose::Bool=false)
DIM=length(μ) # The number of dimensions
if !(DIM==length(σ))
error("length(μ)=$(DIM) while length(σ)=$(length(σ)). The number of mean (μ) points must be equal to the number of standard deviation (σ) points. Please correct and re-try.")
end
if !(DIM==length(mx))
error("length(μ)=$(DIM) while length(mx)=$(length(mx)). The number of upper bound (mx) points must equal the number of mean (μ) points. Please correct and re-try.")
end
if !(DIM==length(mn))
error("length(μ)=$(DIM) while length(mn)=$(length(mx)). The number of lower bound (mn) points must equal the number of mean (μ) points. Please correct and re-try.")
end
if !(DIM==length(n))
verbose && println("Matching length of grid points 'n' to $(DIM)-dimensional space... ")
n = repeat(vcat(n),DIM)
end
verbose && println("Upper bound(s) for $(DIM)-dimensional grid: $(mx)")
verbose && println("Lower bound(s) for $(DIM)-dimensional grid: $(mn)")
verbose && println("Number of grid points (in each dimension): $(n)")
v = σ.^2 # Compute the variance from the standard deviation
verbose && println("Creating $(DIM)-dimensional grid for Gaussian distribution... ")
query_vecs_n_inds = () # A tuple to hold all query points and their indices. Structure: ((vector,indices),(vector,indices),...)
for i in 1:DIM # For all grid dimensions...
query_vecs_n_inds = tuple(query_vecs_n_inds[:]...,collect(zip(collect(range(mn[i],stop=mx[i],length=n[i])),1:n[i]))) # Add the (vector,indices) pairs one by one into a big tuple (tuples are immutable, hence the cumbersome code)
end
query_points_n_coords = Iterators.product(query_vecs_n_inds...) # Create a long list of all reconstruction space grid points and their coordinates by computing a product between all query point-index vectors. Example structure (if 3 dimensions): [((x1_1,1),(x2_1,1),(x3_1,1)),((x1_2,2),(x2_1,1),(x3_1,1)),...]
verbose && print("Computing Gaussian distribution with mean $(μ) and standard deviation $(σ)...")
gauss_distr = zeros(tuple(n...)) # Pre-allocate Gaussian distribution
for query_point_n_coord in query_points_n_coords
point = [p for p in map(x-> x[1],query_point_n_coord)] # The point to compute the Gaussian at. E.g. (100.0,0.3) in energy (keV),pitch
coord = map(x-> x[2],query_point_n_coord) # The coordinate of that point. E.g. (53,14)
gauss_distr[coord...] = ((2*pi)^(-DIM/2))*inv(sqrt(det(diagm(v)))) *exp(-0.5*transpose(point - μ)*inv(diagm(v))*(point - μ))
end
verbose && println("Done!")
if floor_level>0
verbose && print("Grid points with values below $(floor_level)*maximum(gauss_distr) will be manually set to 0.0...")
max_g = maximum(gauss_distr)
gauss_distr = map(x-> x<floor_level*max_g ? 0.0 : x, gauss_distr)
verbose && println("Done!")
end
return gauss_distr
end
"""
erf(x::Real)
erf(x; resolution::Int64 = 1000, sigma=1/sqrt(2))
This is the error function, defined as
erf(x) = (2/(sqrt(2*π)*σ)) ∫ exp(-t^2 / (2 σ^2)) dt
where the lower integration limit is 0 and the upper integration limit is x. σ is set to 1/sqrt(2) by default, via
the keyword argument 'sigma'. This function is a quick approximation, since a sum is used instead of integration.
The number of summation elements can be set via the 'resolution' keyword argument.
"""
function erf(x::Real; resolution::Int64=1000, sigma::Float64=1/sqrt(2))
if x<0.0
t_array = collect(range(x,stop=0.0,length=resolution))
else
t_array = collect(range(0.0,stop=x,length=resolution))
end
dt = x/resolution
return clamp((2/(sqrt(2*pi)*sigma))*sum(dt .*exp.((-1/(2*sigma^2)) .*(t_array).^2)), -1, 1) # Clamp polishes insufficient resolution issues
end
"""
erfc(x::Real)
erfc(x; resolution::Int64 = 1000)
This is the complementary error function, defined as
erfc(x) = 1 - erf(x)
This function is a quick approximation, since a sum is used instead of integration for erf(x).
The number of summation elements can be set via the 'resolution' keyword argument.
"""
function erfc(x::Real; resolution::Int64=1000, sigma::Float64=1/sqrt(2))
return 1 - erf(x; resolution=resolution, sigma=sigma)
end
"""
closest_index(myArray, val)
Finds (the index of) the closest value to val in myArray. Just a function found on the internet as open source. It works.
"""
function closest_index(x::AbstractArray, val::Number)
ibest = first(eachindex(x))
dxbest = abs(x[ibest]-val)
for I in eachindex(x)
dx = abs(x[I]-val)
if dx < dxbest
dxbest = dx
ibest = I
end
end
ibest
end
###### Geometry
"""
rz_grid(rmin, rmax, nr, zmin, zmax, nz, phimin=0.0, phimax=0.0, nphi=1)
Creates an interpolation grid.
#### Input arguments
## rmin - Minimum radius [cm]
## rmax - Maximum radius [cm]
## nr - Number of radii
## zmin - Minimum z value [cm]
## zmax - Maximum z value [cm]
## nz - Number of z values
## phimin - Minimum Phi value [rad]
## phimax - Maximum Phi value [rad]
## nphi - Number of Phi values
#### Return Value
## Interpolation grid dictionary
####Example Usage
## ```julia
## julia> grid = rz_grid(0,200.0,200,-100,100,200,phimin=4*np.pi/3,phimax=5*np.pi/3,nphi=5)
## ```
"""
function rz_grid(rmin::Union{Int64,Float64}, rmax::Union{Int64,Float64}, nr::Int64, zmin::Union{Int64,Float64}, zmax::Union{Int64,Float64}, nz::Int64; phimin::Union{Int64,Float64}=0.0, phimax::Union{Int64,Float64}=0.0, nphi::Int64=1)
dr = (rmax - rmin) / nr
dz = (zmax - zmin) / nz
dphi = (phimax - phimin) / nphi
r = rmin .+ dr * (0:(nr-1))
z = zmin .+ dz * (0:(nz-1))
phi = phimin .+ dphi * (0:(nphi-1))
r2d = repeat(r,1,nz)
z2d = repeat(z,1,nr)'
return grid(r2d,z2d,r,z,phi,nr,nz,nphi)
end
###### Particle species functions
"""
getGCP(species_identifier::AbstractString)
The same function as is provided by OWCF/misc/species_func.jl. However, since it's needed in certain functions below, which are also
a part of extra/dependencies.jl, it needs to be defined here as well. Sometimes, it's more efficient to only load OWCF/misc/species_func.jl,
since one does not always want all the packages needed for dependencies.jl. Therefore, the functions below are also loadable as a separate
script, i.e. OWCF/misc/species_func.jl.
"""
function getGCP(species_identifier::AbstractString)
if lowercase(species_identifier)=="d"
return GuidingCenterOrbits.GCDeuteron
elseif lowercase(species_identifier)=="t"
return GuidingCenterOrbits.GCTriton
elseif lowercase(species_identifier)=="3he"
return GuidingCenterOrbits.GCHelium3
elseif lowercase(species_identifier)=="p"
return GuidingCenterOrbits.GCProton
elseif lowercase(species_identifier)=="e"
return GuidingCenterOrbits.GCElectron
elseif (lowercase(species_identifier)=="alpha" || lowercase(species_identifier)=="4he")
return GuidingCenterOrbits.GCAlpha
else
error("getSpeciesMass got unknown species as input. Please re-try.")
end
end
function getSpeciesMass(species_identifier::AbstractString)
return (getGCP(species_identifier)(0.0,0.0,0.0,0.0)).m # 0.0, 0.0, 0.0, 0.0 is just to activate the function
end
function getSpeciesAmu(species_identifier::AbstractString)
return getSpeciesMass(species_identifier) / GuidingCenterOrbits.mass_u
end
function getSpeciesEcu(species_identifier::AbstractString)
return (getGCP(species_identifier)(0.0,0.0,0.0,0.0)).q # 0.0, 0.0, 0.0, 0.0 is just to activate the function
end
function getSpeciesCharge(species_identifier::AbstractString)
return getSpeciesEcu(species_identifier) * GuidingCenterOrbits.e0
end
###### Basic physics
"""
debye_length(n_e::Real, T_e::Real, species_th_vec::Vector{String}, n_th:vec::Vector{T}, T_th_vec::Vector{T}) where T<:Real
debye_length(-||-; species_2::String=species_1, n_2::Real=n_1, T_2::Real=T_1)
Compute the (exact) plasma debye length. The inputs are:
- n_e: The electron density [m^⁻3]
- T_e: The electron temperature [keV]
- species_th_vec: A vector containing all thermal species string identifiers, e.g. ["D", "T", "3he"] [-]
- n_th_vec: A vector containing all thermal species densities. n_th_vec[i] is the density of species_th_vec[i] [m^-3]
- T_th_vec: A vector containing all thermal species temperatures. T_th_vec[i] is the temperature of species_th_vec[i] [keV]
The outputs are:
- debye_length: The debye length [m]
"""
function debye_length(n_e::Float64, T_e::Real, species_th_vec::Vector{String}, n_th_vec::Vector{T} where {T<:Real}, T_th_vec::Vector{T} where {T<:Real})
temp_e = T_e*1000*(GuidingCenterOrbits.e0)/kB
temp_th_vec = (1000*(GuidingCenterOrbits.e0)/kB) .*T_th_vec
Z_th_vec = getSpeciesEcu.(species_th_vec)
denom = (n_e/temp_e) + reduce(+, Z_th_vec .*Z_th_vec .*n_th_vec ./temp_th_vec)
return sqrt((ϵ0*kB/(GuidingCenterOrbits.e0)^2)/denom)
end
"""
gyro_radius(M::AbstractEquilibrium,p::GCParticle)
Compute the (relativistic) gyro radius for guiding-center particle p, given the magnetic
equilibrium M. Output in meters. Take relativity into account. The inputs are:
- M: An abstract equilibrium. Either from an .eqdsk file or an output file from OWCF/extra/compSolovev.jl [-]
- p: The guiding-center particle object for the particle species of interest, e.g. GCDeuteron [-]
The output is:
- gyro_radius: The relativistically correct gyro radius [m]
"""
function gyro_radius(M::AbstractEquilibrium,p::GCParticle)
γ = GuidingCenterOrbits.lorentz_factor(p)
Babs = norm(Bfield(M, p.r, p.z)) # Magnetic field magnitude. Tesla
m = p.m # Mass of particle
KE = (γ-1)*p.energy # Relativistic kinetic energy. keV
mc2 = m*GuidingCenterOrbits.c0*GuidingCenterOrbits.c0 # Rest energy. Joule
KE_j = GuidingCenterOrbits.e0*KE*1e3 # Kinetic energy. Joule
p_rel2 = ((KE_j + mc2)^2 - mc2^2)/(GuidingCenterOrbits.c0*GuidingCenterOrbits.c0) # Relativistic momentum
p_perp2 = p_rel2*(1-p.pitch^2) # Square of relativistic perpendicular momentum
return sqrt(p_perp2) / (abs(GuidingCenterOrbits.e0*p.q)*Babs*γ)
end
"""
spitzer_slowdown_time(n_e, T_e, species_f, species_th_vec, n_th_vec, T_th_vec)
spitzer_slowdown_time(-||-; plasma_model = :texas, returnExtra = false)
Compute the non-relativistic Spitzer slowing-down time (in seconds) for fast-ion species 'species_f', following the equation in the ITER Physics Basis (http://sites.apam.columbia.edu/fusion/IPB_Chap_5.pdf).
Assume multiple thermal species via the vector inputs. By default, use the texas (University of Texas) model for the Coloumb logarithm.
if returnExtra, in addition to τ_s, return the coulomb logarithm as well as the Debye length. The inputs are:
- n_e: The electron density [m^-3]
- T_e: The electron temperature [keV]
- species_f: The beam injection particle species, e.g. "D", "T" etc [-]
- species_th_vec: A vector containing all thermal species string identifiers, e.g. ["D", "T", "3he"] [-]
- n_th_vec: A vector containing all thermal species densities. n_th_vec[i] is the density of species_th_vec[i] [m^-3]
- T_th_vec: A vector containing all thermal species temperatures. T_th_vec[i] is the temperature of species_th_vec[i] [keV]
The keyword arguments are:
- plasma_model: The model to use for the plasma parameter. Currently supported :salewski or :texas [-]
- returnExtra: If true, in addition to the Spitzer slowing-down time, the Coulomb logarithm and Debye length will be returned as well [-]
"""
function spitzer_slowdown_time(n_e::Real, T_e::Real, species_f::String, species_th_vec::Union{String,Vector{String}}, n_th_vec::Union{T,Vector{T}} where {T<:Real}, T_th_vec::Union{T,Vector{T}} where {T<:Real}; plasma_model::Symbol = :texas, returnExtra::Bool = false)
species_th_vec = vcat(species_th_vec) # In case input was a String, and not a Vector
n_th_vec = vcat(n_th_vec) # -||-
T_th_vec = vcat(T_th_vec) # -||-
ϵ0 = 8.8541878128e-12 # Permittivity of free space
m_f = getSpeciesMass(species_f) # The fast-ion species mass, kg
m_e = (GuidingCenterOrbits.e_amu)*(GuidingCenterOrbits.mass_u) # Electron mass, kg
λ_D = debye_length(n_e, T_e, species_th_vec, n_th_vec, T_th_vec)
if plasma_model==:salewski
Λ = 6*pi*n_e*λ_D^3 # From M. Salewski, A.H. Nielsen, Plasma Physics: lectures notes, 2021.
Λ_c = Λ # Actually, this is Λ_c ≈ Λ, where Λ is the plasma parameter.
elseif plasma_model==:texas
Z_th_vec = getSpeciesEcu.(species_th_vec) # Atomic charge number for all thermal plasma species
q_th_vec = (GuidingCenterOrbits.e0) .*Z_th_vec # Charge (in Coulombs) for all thermal plasma species
q2_avg = reduce(+,map(x-> x[1]*x[2],collect(Iterators.product(q_th_vec,q_th_vec))[:]))/(length(q_th_vec)^2) # The average of q_i*q_j for all species pairs (i,j)
r_closest = q2_avg/(4*pi*ϵ0*T_e*GuidingCenterOrbits.e0*1000) # r_closest = <q_i*q_j>/(4*pi*ϵ0*T). Mean value of closest approach, University of Texas. Assume same termperature.
Λ_c = λ_D/r_closest
else
error("Currently supported models for the plasma parameter are :salewski and :texas. Please correct and re-try.")
end
coulomb_log = log(Λ_c)
A_D = (n_e*((GuidingCenterOrbits.e0)^4)*coulomb_log)/(2*pi*(ϵ0^2)*(m_f^2))
τ_s = (3*sqrt(2*pi)*((T_e*GuidingCenterOrbits.e0*1000)^(3/2)))/(sqrt(m_e)*m_f*A_D)
if returnExtra
return τ_s, coulomb_log, λ_D
end
return τ_s
end
###### Data manipulation
"""
add_noise(S, b)
add_noise(-||-; k=0.1)
A function that adds noise to a signal in a consistent way. The function adds both background noise as well as
noise to the signal itself. The level of the background noise and the signal noise is determined by the b
and k variables respectively. The levels are input as a fraction of the mean of the original signal S strength.
"""
function add_noise(s::AbstractVector, b::Union{Float64, AbstractVector}; k::Float64=0.1)
sn = max.(s,0.0) .+ k.*(mean(sqrt.(abs.(s)))).*rand.(Normal.(0.0, max.(sqrt.(max.(s,0)), sqrt.(b))))
err = k.*mean(sqrt.(abs.(s))).*max.(sqrt.(max.(s,0)), sqrt.(b))
return sn, err
end
"""
estimate_noise(S_tot)
A function that estimates the noise of a signal S. With the 'bad_inds' keyword argument,
estimate the noise of the S_tot-elements with the specific array indices in 'bad_inds'.
Assume that the true signal 'S_tot' is a smooth function without the noise.
"""
function estimate_noise(S_tot::AbstractVector; bad_inds::Union{Nothing,AbstractVector}=nothing)
noise_tot = zeros(length(S_tot))
if bad_inds==nothing
bad_inds = 1:length(S_tot)
end
for bad_ind in bad_inds # A simple noise estimation algorithm
if bad_ind==1 # If the bad index is the very first element, we need to extrapolate twice
S_bad_one_ahead_mean = (S_tot[bad_ind]+S_tot[bad_ind+1]+S_tot[bad_ind+2])/3
S_extrap = 2*S_tot[bad_ind] - S_bad_one_ahead_mean
S_worse_mean = (S_extrap+S_tot[bad_ind]+S_tot[bad_ind+1])/3
S_extrap_extrap = 2*S_extrap - S_worse_mean
S_bad_mean = (S_extrap_extrap+S_extrap+S_tot[bad_ind]+S_tot[bad_ind+1]+S_tot[bad_ind+2])/5 # The average of all five points close to the point of interest
elseif bad_ind==2 # If the bad index is the second element, we need to extrapolate once
S_worse_mean = (S_tot[bad_ind-1]+S_tot[bad_ind]+S_tot[bad_ind+1])/3
S_extrap = 2*S_tot[bad_ind-1] - S_worse_mean
S_bad_mean = (S_extrap+S_tot[bad_ind-1]+S_tot[bad_ind]+S_tot[bad_ind+1]+S_tot[bad_ind+2])/5 # The average of all five points close to the point of interest
elseif bad_ind==(length(S_tot)-1) # If the bad index is the next to last element, we need to extrapolate once
S_worse_mean = (S_tot[bad_ind-1]+S_tot[bad_ind]+S_tot[bad_ind+1])/3
S_extrap = 2*S_tot[bad_ind+1]-S_worse_mean
S_bad_mean = (S_tot[bad_ind-2]+S_tot[bad_ind-1]+S_tot[bad_ind]+S_tot[bad_ind+1]+S_extrap)/5 # The average of all five points close to the point of interest
elseif bad_ind==length(S_tot) # If the bad index is the very last element, we need to extrapolate twice
S_bad_one_behind_mean = (S_tot[bad_ind-2]+S_tot[bad_ind-1]+S_tot[bad_ind])/3
S_extrap = 2*S_tot[bad_ind]-S_bad_one_behind_mean
S_worse_mean = (S_tot[bad_ind-1]+S_tot[bad_ind]+S_extrap)/3
S_extrap_extrap = 2*S_extrap - S_worse_mean
S_bad_mean = (S_tot[bad_ind-2]+S_tot[bad_ind-1]+S_tot[bad_ind]+S_extrap+S_extrap_extrap)/5 # The average of all five points close to the point of interest
else
S_bad_mean = (S_tot[bad_ind-2]+S_tot[bad_ind-1]+S_tot[bad_ind]+S_tot[bad_ind+1]+S_tot[bad_ind+2])/5 # The average of all five points close to the point of interest
end
noise_tot[bad_ind] = abs(S_tot[bad_ind] - S_bad_mean) # Noise (or error)
end
return noise_tot
end
"""
apply_instrumental_response(S, Ed_array, instrumental_response_input, instrumental_response_output, instrumental_response_matrix)
apply_instrumental_response(-||-; lo=nothing, hi=nothing, instrumental_response=true)
Apply instrumental response to the signal S via the matrix instrumental_response matrix. The input and output axes of the instrumental response
matrix are specified as vector instrumental_response_input and instrumental_response_output. The index of the element closest to minimum(Ed_array)
in instrumental_response_input can be specified via the lo keyword argument. The index of the element closest to maximum(Ed_array) in instrumental_response_input
can be specified via the hi keyword argument.
"""
function apply_instrumental_response(S::AbstractVector, Ed_array::AbstractVector, instrumental_response_input::AbstractVector, instrumental_response_output::AbstractVector, instrumental_response_matrix::AbstractMatrix; lo::Union{Int64,Nothing}=nothing, hi::Union{Int64,Nothing}=nothing)
if isnothing(lo) # If no lower bound for the instrumental response input has been provided...
lo = findfirst(x-> x>minimum(Ed_array),instrumental_response_input) # Find it
end
if isnothing(hi) # Similar as lo, but hi
hi = findlast(x-> x<maximum(Ed_array),instrumental_response_input)
end
if isnothing(lo) || isnothing(hi) # If extrema(Ed_array) are (partially) outside of extrema(instrumental_response_input)...
@warn "Instrumental response matrix input completely outside weight function measurement bin range. No diagnostic response will be applied."
return S
else
if lo==1
@warn "Lower bound of instrumental response matrix input might not be low enough to cover weight function measurement bin range. Diagnostic response representation might be inaccurate."
end
if hi==length(instrumental_response_input)
@warn "Upper bound of instrumental response matrix input might not be high enough to cover weight function measurement bin range. Diagnostic response representation might be inaccurate."
end
instrumental_response_matrix = (instrumental_response_matrix[lo:hi,:])' # Take the transpose, to go from (input, output) shape to (output, input) shape
itp = LinearInterpolation(Ed_array,S) # Create interpolation object for S
S_itp = itp.(instrumental_response_input[lo:hi]) # Find interpolation values for S and instrumental_response_input points
S_out = instrumental_response_matrix * S_itp # The instrumental response is applied
return S_out
end
end
"""
apply_instrumental_response(Q, Ed_array, instrumental_response_input, instrumental_response_output, instrumental_response_matrix)
apply_instrumental_response(-||-; lo=nothing, hi=nothing, instrumental_response=true)
Apply the instrumental response given by instrumental_response_matrix to each column of Q. Assume the diagnostic spectrum is accessed
via the first dimension of Q. That is, Q[1,...] corresponds to the first weight function, Q[2,...] corresponds to the second weight function
and so on.
"""
function apply_instrumental_response(Q::Array{Float64,2}, Ed_array::AbstractVector, instrumental_response_input::AbstractVector, instrumental_response_output::AbstractVector, instrumental_response_matrix::AbstractMatrix; lo::Union{Int64,Nothing}=nothing, hi::Union{Int64,Nothing}=nothing)
if isnothing(lo)
lo = findfirst(x-> x>minimum(Ed_array),instrumental_response_input)
end
if isnothing(hi)
hi = findlast(x-> x<maximum(Ed_array),instrumental_response_input)
end
if isnothing(lo) || isnothing(hi)
@warn "Instrumental response matrix input completely outside weight function measurement bin range. No diagnostic response will be applied."
return Q
else
Q_out = zeros(length(instrumental_response_output), size(Q,2))
if lo==1
@warn "Lower bound of instrumental response matrix input might not be low enough to cover weight function measurement bin range. Diagnostic response representation might be inaccurate."
end
if hi==length(instrumental_response_input)
@warn "Upper bound of instrumental response matrix input might not be high enough to cover weight function measurement bin range. Diagnostic response representation might be inaccurate."
end
for i=1:size(Q,2) # Apply instrumental response to each column
Q_out[:,i] = apply_instrumental_response(Q[:,i], Ed_array, instrumental_response_input, instrumental_response_output, instrumental_response_matrix; lo=lo, hi=hi)
end
return Q_out
end
end
"""
apply_instrumental_response(Q, Ed_array, instrumental_response_input, instrumental_response_output, instrumental_response_matrix)
apply_instrumental_response(-||-; lo=nothing, hi=nothing, instrumental_response=true)
Apply the instrumental response given by instrumental_response_matrix to each matrix of Q. Assume the diagnostic spectrum is accessed
via the first dimension of Q. That is, Q[1,...] corresponds to the first weight function, Q[2,...] corresponds to the second weight function
and so on.
"""
function apply_instrumental_response(Q::Array{Float64,3}, Ed_array::AbstractVector, instrumental_response_input::AbstractVector, instrumental_response_output::AbstractVector, instrumental_response_matrix::AbstractMatrix; lo::Union{Int64,Nothing}=nothing, hi::Union{Int64,Nothing}=nothing)
if isnothing(lo)
lo = findfirst(x-> x>minimum(Ed_array),instrumental_response_input)
end
if isnothing(hi)
hi = findlast(x-> x<maximum(Ed_array),instrumental_response_input)
end
if isnothing(lo) || isnothing(hi)
@warn "Instrumental response matrix input completely outside weight function measurement bin range. No diagnostic response will be applied."
return Q
else
Q_out = zeros(length(instrumental_response_output), size(Q,2), size(Q,3))
if lo==1
@warn "Lower bound of instrumental response matrix input might not be low enough to cover weight function measurement bin range. Diagnostic response representation might be inaccurate."
end
if hi==length(instrumental_response_input)
@warn "Upper bound of instrumental response matrix input might not be high enough to cover weight function measurement bin range. Diagnostic response representation might be inaccurate."
end
for i=1:size(Q,3) # Apply instrumental response to each matrix
Q_out[:,:,i] = apply_instrumental_response(Q[:,:,i], Ed_array, instrumental_response_input, instrumental_response_output, instrumental_response_matrix; lo=lo, hi=hi)
end
return Q_out
end
end
"""
apply_instrumental_response(Q, Ed_array, instrumental_response_input, instrumental_response_output, instrumental_response_matrix)
apply_instrumental_response(-||-; lo=nothing, hi=nothing, instrumental_response=true)
Apply the instrumental response given by instrumental_response_matrix to each 3D array of Q. Assume the diagnostic spectrum is accessed
via the first dimension of Q. That is, Q[1,...] corresponds to the first weight function, Q[2,...] corresponds to the second weight function
and so on.
"""
function apply_instrumental_response(Q::Array{Float64,4}, Ed_array::AbstractVector, instrumental_response_input::AbstractVector, instrumental_response_output::AbstractVector, instrumental_response_matrix::AbstractMatrix; lo::Union{Int64,Nothing}=nothing, hi::Union{Int64,Nothing}=nothing)
if isnothing(lo)
lo = findfirst(x-> x>minimum(Ed_array),instrumental_response_input)
end
if isnothing(hi)
hi = findlast(x-> x<maximum(Ed_array),instrumental_response_input)
end
if isnothing(lo) || isnothing(hi)
@warn "Instrumental response matrix input completely outside weight function measurement bin range. No diagnostic response will be applied."
return Q
else
Q_out = zeros(length(instrumental_response_output), size(Q,2), size(Q,3), size(Q,4))
if lo==1
@warn "Lower bound of instrumental response matrix input might not be low enough to cover weight function measurement bin range. Diagnostic response representation might be inaccurate."
end
if hi==length(instrumental_response_input)
@warn "Upper bound of instrumental response matrix input might not be high enough to cover weight function measurement bin range. Diagnostic response representation might be inaccurate."
end
for i=1:size(Q,4) # Apply instrumental response to each 3D array
Q_out[:,:,:,i] = apply_instrumental_response(Q[:,:,:,i], Ed_array, instrumental_response_input, instrumental_response_output, instrumental_response_matrix; lo=lo, hi=hi)
end
return Q_out
end
end
"""
apply_instrumental_response(Q, Ed_array, instrumental_response_input, instrumental_response_output, instrumental_response_matrix)
apply_instrumental_response(-||-; lo=nothing, hi=nothing)
Apply the instrumental response given by instrumental_response_matrix to each 4D array of Q. Assume the diagnostic spectrum is accessed
via the first dimension of Q. That is, Q[1,...] corresponds to the first weight function, Q[2,...] corresponds to the second weight function
and so on.
"""
function apply_instrumental_response(Q::Array{Float64,5}, Ed_array::AbstractVector, instrumental_response_input::AbstractVector, instrumental_response_output::AbstractVector, instrumental_response_matrix::AbstractMatrix; lo::Union{Int64,Nothing}=nothing, hi::Union{Int64,Nothing}=nothing)
if isnothing(lo)
lo = findfirst(x-> x>minimum(Ed_array),instrumental_response_input)
end
if isnothing(hi)
hi = findlast(x-> x<maximum(Ed_array),instrumental_response_input)
end
if isnothing(lo) || isnothing(hi)
@warn "Instrumental response matrix input completely outside weight function measurement bin range. No diagnostic response will be applied."
return Q
else
Q_out = zeros(length(instrumental_response_output), size(Q,2), size(Q,3), size(Q,4), size(Q,5))
if lo==1
@warn "Lower bound of instrumental response matrix input might not be low enough to cover weight function measurement bin range. Diagnostic response representation might be inaccurate."
end
if hi==length(instrumental_response_input)
@warn "Upper bound of instrumental response matrix input might not be high enough to cover weight function measurement bin range. Diagnostic response representation might be inaccurate."
end
for i=1:size(Q,5) # Apply instrumental response to each 4D array
Q_out[:,:,:,:,i] = apply_instrumental_response(Q[:,:,:,:,i], Ed_array, instrumental_response_input, instrumental_response_output, instrumental_response_matrix; lo=lo, hi=hi)
end
return Q_out
end
end
###### (Drift) Orbit-related functions
"""
orbit_grid(M,false,eo,po,ro)
orbit_grid(-||-, q=1, amu=OrbitTomography.H2_amu, kwargs... )
This function is just like OrbitTomography.orbit_grid, but allows execution without progress bar.
Good for HPC batch job submission, to not overrun log files.
"""
function orbit_grid(M::AbstractEquilibrium, visualizeProgress::Bool, eo::AbstractVector, po::AbstractVector, ro::AbstractVector;
q::Int64 = 1, amu = OrbitTomography.H2_amu, kwargs...)
nenergy = length(eo)
npitch = length(po)
nr = length(ro)
orbit_index = zeros(Int,nenergy,npitch,nr)
class = fill(:incomplete,(nenergy,npitch,nr))
tau_t = zeros(Float64,nenergy,npitch,nr)
tau_p = zeros(Float64,nenergy,npitch,nr)
norbs = nenergy*npitch*nr
subs = CartesianIndices((nenergy,npitch,nr))
if visualizeProgress
prog = ProgressMeter.Progress(norbs)
channel = RemoteChannel(()->Channel{Bool}(norbs), 1)
orbs = fetch(@sync begin
@async while take!(channel)
ProgressMeter.next!(prog)
end
@async begin
orbs = @distributed (vcat) for i=1:norbs
ie,ip,ir = Tuple(subs[i])
c = GuidingCenterOrbits.EPRCoordinate(M,eo[ie],po[ip],ro[ir],q=q,amu=amu)
try
o = GuidingCenterOrbits.get_orbit(M, c; kwargs...)
catch
o = GuidingCenterOrbits.Orbit(EPRCoordinate(;q=q,amu=amu),:incomplete)
end
if o.class in (:incomplete,:invalid,:lost)
o = Orbit(o.coordinate,:incomplete)
end
put!(channel, true)
o
end
put!(channel, false)
orbs
end
end)
else
orbs = @distributed (vcat) for i=1:norbs
ie,ip,ir = Tuple(subs[i])
c = GuidingCenterOrbits.EPRCoordinate(M,eo[ie],po[ip],ro[ir],q=q,amu=amu)
try
o = GuidingCenterOrbits.get_orbit(M, c; kwargs...)
catch
o = GuidingCenterOrbits.Orbit(EPRCoordinate(;q=q,amu=amu),:incomplete)
end
if o.class in (:incomplete,:invalid,:lost)
o = Orbit(o.coordinate,:incomplete)
end
o
end
end
for i=1:norbs
class[subs[i]] = orbs[i].class
tau_p[subs[i]] = orbs[i].tau_p
tau_t[subs[i]] = orbs[i].tau_t
end
grid_index = filter(i -> orbs[i].class != :incomplete, 1:norbs)
orbs = filter(x -> x.class != :incomplete, orbs)
norbs = length(orbs)
orbit_index[grid_index] = 1:norbs
return orbs, OrbitTomography.OrbitGrid(eo,po,ro,fill(1,norbs),orbit_index,class,tau_p,tau_t)
end
"""
mu_func(energy, B, Pϕ, Ψ, RBϕ)
mu_func(-||-; m=GuidingCenterOrbits.H2_amu*GuidingCenterOrbits.mass_u, q=1*GuidingCenterOrbits.e0)
Compute the magnetic moment mu, given the fast-ion energy, magnetic field B, toroidal canonical angular momentum Pϕ, magnetic flux Ψ and flux function RBϕ.
Use a 0-cap, meaning that all negative values are set to 0.
"""
function mu_func(energy::T, B::T, Pϕ::T, Ψ::T, RBϕ::T;m::Float64=GuidingCenterOrbits.H2_amu*GuidingCenterOrbits.mass_u, q::Float64=1*GuidingCenterOrbits.e0) where {T}
res = energy/B - (B/(2*m)) * ((Pϕ-q*Ψ)/RBϕ)^2
return (res > zero(energy)) ? res : zero(energy)
end
"""
get_orbel_volume(myOrbitGrid, os_equidistant)
Get the orbit element volume of this specific orbit-grid.
If os_equidistant=false, then return a 3D array with all the orbit volumes.
Otherwise, just a scalar.
"""
function get_orbel_volume(og::OrbitGrid, os_equidistant::Bool)
if os_equidistant
dRm = abs(og.r[2]-og.r[1]) # The difference between the first and second element should be representative for the whole grid, due to equidistancy
dE = abs(og.energy[2]-og.energy[1])
dpm = abs(og.pitch[2]-og.pitch[1])
dO = dE*dpm*dRm
return dO
else
dO = zeros(length(og.energy), length(og.pitch), length(og.r)) # If not equidistant, create a 3D array with zeros
for Ei=1:length(og.energy) # For all energies...
if Ei==length(og.energy) # If at the edge of the orbit grid...
# Assume edge orbit-element volume to be same as next-to-edge
dEi = abs(og.energy[end]-og.energy[end-1])
else
dEi = abs(og.energy[Ei+1]-og.energy[Ei])
end
for pmi=1:length(og.pitch)
if pmi==length(og.pitch)
dpmi = abs(og.pitch[end]-og.pitch[end-1])
else
dpmi = abs(og.pitch[pmi+1]-og.pitch[pmi])
end
for Rmi=1:length(og.r)
if Rmi==length(og.r)
dRmi = abs(og.r[end]-og.r[end-1])
else
dRmi = abs(og.r[Rmi+1]-og.r[Rmi])
end
dO[Ei, pmi, Rmi] = dEi*dpmi*dRmi # Replace the zero-element in the 3D array with the resulting orbit-element volume
end
end
end
return dO
end
end
###### UNLABELLED FUNCTIONS BELOW
"""
interpFps(f,E,p,R,z,Eq,pq,Rq,zq)
The 4D particle-space fast-ion distribution f, with the corresponding
grid vectors in energy E, pitch p, radius R and vertical position z, are supplied as input. Evaluate the
fast-ion distribution at the query points (Eq, pq, Rq and zq) using linear interpolation in 4D.
Return the resulting 4D fast-ion distribution.
"""
function interpFps(f::AbstractArray, E::AbstractArray, p::AbstractArray, R::AbstractArray, z::AbstractArray, Eq::AbstractArray, pq::AbstractArray, Rq::AbstractArray, zq::AbstractArray; debug::Bool=false)
# Assume equidistant elements
nodes = (E, p, R, z)
# Create interpolation object
itp = interpolate(nodes, f, Gridded(Linear()))
fq = zeros(length(Eq),length(pq),length(Rq),length(zq)) # Pre-allocate 4D array
numObadInds = 0
for Ei in eachindex(Eq)
for pi in eachindex(pq)
for Ri in eachindex(Rq)
for zi in eachindex(zq)
try
fq[Ei,pi,Ri,zi] = itp(Eq[Ei],pq[pi],Rq[Ri],zq[zi]) # Interpolate at 4D query point
catch
numObadInds += 1
debug && println("(Ei: $(Ei), pi: $(pi), Ri: $(Ri), zi: $(zi)) <--- Interpolation failed for this index") # Print if failed (should not happen)
end
end
end
end
end
if numObadInds > 0
perc_frac = (numObadInds / length(fq))*100
@warn "The interpolation algorithm failed to interpolate $(numObadInds) (E,p,R,z) points ($(round(perc_frac,sigdigits=3)) %)"
end
return fq
end
"""
do4Dto2D(og, W)
do4Dto2D(-||-, returnComplement=false)
This function converts orbit weights of the form (channel,E,pm,Rm) to the form
(channel,orbits). 'orbits' consists of the valid orbits for the (E,pm,Rm)-grid.
Usually, in orbit space, for a given orbit grid only about 67% of the (E,pm,Rm)-points
corresponds to valid orbits. So extracting the actual valid points as a 1D vector makes sense.
If returnComplement, then a 1D array with all the weight values
that should be zero is returned (C), with the corresponding coordinates (Ccoords) and indices (Cindices) as arrays of triplets.
The sum of the C array should be zero. If it is not, then invalid orbits have been given a non-zero weight. Good for debugging.
"""
function do4Dto2D(og::OrbitGrid, W4D::AbstractArray; returnComplement::Bool=false)
W = zeros(size(W4D,1),length(og.counts)) # Pre-allocate the 2D (channel,orbits) matrix
if returnComplement
C = zeros(length(W4D)-length(og.counts)) #Pre-allocate
Ccoords = Array{Tuple{Float64,Float64,Float64},1}(undef,length(W4D)-length(og.counts)) #Pre-allocate
Cindices = Array{Tuple{Int64,Int64,Int64},1}(undef,length(W4D)-length(og.counts)) #Pre-allocate
ii = 1
end
for c in axes(W4D,1)
for Ei in axes(W4D,2)
for pmi in axes(W4D,3)
for Rmi in axes(W4D,4)
o_index = og.orbit_index[Ei,pmi,Rmi] # By the method of L. Stagner, every valid orbit corresponds to a non-zero integer: their index. The invalid and lost orbits are simply represented by a zero as their index.
if o_index==0 && returnComplement
C[ii] = W4D[c,Ei,pmi,Rmi]
Ccoords[ii] = (og.energy[Ei],og.pitch[pmi],og.r[Rmi])
Cindices[ii] = (Ei,pmi,Rmi)
ii = ii + 1
elseif o_index==0
else
W[c,o_index] = W4D[c,Ei,pmi,Rmi] # Put the orbit weight at the correct orbit position (column) in the 2D weight matrix
end
end
end
end
end
if returnComplement
return W,C,Ccoords,Cindices
else
return W
end
end
"""
flip_along_pm(W)
flip_along_pm(W, with_channels=false)
Flip the elements of an orbit-space function along the pm axis. If with_channels=true,
then a 4D orbit-space function on the form (channel,E,pm,Rm) or (channel,Rm,pm,E) is assumed.
Otherwise a form of (E,pm,Rm) is assumed.
"""
function flip_along_pm(W::AbstractArray; with_channels::Bool=false)
Wres = zeros(size(W))
if with_channels
for pmi in axes(W,3)
Wres[:,:,(end+1)-pmi,:] .= W[:,:,pmi,:]
end
else
for pmi in axes(W,2)
Wres[:,(end+1)-pmi,:] .= W[:,pmi,:]
end
end
return Wres
end
"""
h5to4D(filepath_distr)
h5to4D(filepath_distr, backwards=true, verbose=false)
Load and read a .h5/.hdf5 file containing the data necessary to construct a 4D fast-ion distribution, with dimensions (E,p,R,z).
Automatically assume that the data will be loaded backwards, because that seems to be the case when exporting 4D
data from Python via .h5/.hdf5 files. Returns f, E, p, R, z.
"""
function h5to4D(filepath_distr::AbstractString; backwards::Bool = true, verbose::Bool = false)
if verbose
println("Loading the 4D distribution from .hdf5 file... ")
end
myfile = h5open(filepath_distr,"r")
if haskey(myfile,"F_ps")
verbose && println("Found F_ps key in .hdf5 file.")
fbackwards = read(myfile["F_ps"])
elseif haskey(myfile,"f")
verbose && println("Found f key in .hdf5 file.")
fbackwards = read(myfile["f"])
elseif haskey(myfile,"F_EpRz")
verbose && println("Found F_EpRz key in .hdf5 file.")
fbackwards = read(myfile["F_EpRz"])
else
error("Fast-ion distribution .hdf5 file did not have any expected keys for the distribution (F_ps, f or F_EpRz).")
end
if haskey(myfile,"energy")
energy = read(myfile["energy"])
elseif haskey(myfile,"E_array")
energy = read(myfile["E_array"])
else
error("Fast-ion distribution .hdf5 file did not have any expected keys for the energy grid points (energy or E_array).")
end
if haskey(myfile,"pitch")
pitch = read(myfile["pitch"])
elseif haskey(myfile,"p_array")
pitch = read(myfile["p_array"])
else
error("Fast-ion distribution .hdf5 file did not have any expected keys for the energy grid points (pitch or p_array).")
end
if haskey(myfile,"R")
R = read(myfile["R"])
elseif haskey(myfile,"R_array")
R = read(myfile["R_array"])
else
error("Fast-ion distribution .hdf5 file did not have any expected keys for the R grid points (R or R_array).")
end
if haskey(myfile,"z")
verbose && println("Found z key in .hdf5 file.")
z = read(myfile["z"])
elseif haskey(myfile,"Z")
verbose && println("Found Z key in .hdf5 file.")
z = read(myfile["Z"])
elseif haskey(myfile,"z_array")
z = read(myfile["z_array"])
elseif haskey(myfile,"Z_array")
z = read(myfile["Z_array"])
else
error("Fast-ion distribution .hdf5 file did not have any expected keys for the z grid points (z, Z, z_array or Z_array).")
end
close(myfile)
if (size(fbackwards,1)==length(energy)) && (size(fbackwards,2)==length(pitch)) && (size(fbackwards,3)==length(R)) && (size(fbackwards,4)==length(z))
backwards = false
end
if backwards
verbose && println("Fast-ion distribution data is permutated. Attemping to fix... ")
f = zeros(size(fbackwards,4),size(fbackwards,3),size(fbackwards,2),size(fbackwards,1))
for i in axes(f,1)
for j in axes(f,2)
for k in axes(f,3)
for l in axes(f,4)
f[i,j,k,l] = fbackwards[l,k,j,i]
end
end
end
end
else
f = fbackwards
end
return f, energy, pitch, R, z
end
"""
jld2tohdf5(filepath_jld2)
jld2tohdf5(filepath_jld2; verbose=false)
Convert a .jld2 file to a .hdf5 file.
"""
function jld2tohdf5(filepath_jld2::String; verbose::Bool=false)
filepath_hdf5 = reduce(*, split(filepath_jld2,".")[1:end-1])*".hdf5"
myfile_jld2 = jldopen(filepath_jld2,false,false,false,IOStream)
myfile_hdf5 = h5open(filepath_hdf5,"w")
for key in keys(myfile_jld2)
verbose && println("Writing data: "*key)
data = myfile_jld2[key]
verbose && print(" ↳ Type: $(typeof(data))")
if typeof(data) <: Dict # If the data is a dictionary
verbose && println(" Not ok!")
@warn "Part of the data in "*filepath_jld2*" is a dictionary. The OWCF currently does not support saving dictionaries in .hdf5 file format. The "*key*" data has therefore been omitted."
elseif typeof(data) <: StepRangeLen
verbose && println(" Ok! (Type StepRangeLen converted to type Array via collect())")
data = collect(data)
write(myfile_hdf5,key,data)
else
verbose && println(" Ok!")
write(myfile_hdf5,key,data)
end
end
close(myfile_jld2)
close(myfile_hdf5)
verbose && println("The .jld2 file has been re-written as a .hdf5 file at "*filepath_hdf5)
return filepath_hdf5
end
"""
JLD2to4D(filepath_distr)
JLD2to4D(filepath_distr; verbose=false) # default
Load the fast-ion distribution from a .jld2 file. Assume certain keys to access data. Print information if 'verbose'.
"""
function JLD2to4D(filepath_distr::String; verbose::Bool=false)
myfile = jldopen(filepath_distr,false,false,false,IOStream)
if haskey(myfile,"F_ps")
verbose && println("Found F_ps key in .jld2 file.")
F_EpRz = myfile["F_ps"]
elseif haskey(myfile,"f")
verbose && println("Found f key in .jld2 file.")
F_EpRz = myfile["f"]
elseif haskey(myfile,"F_EpRz")
verbose && println("Found F_EpRz key in .jld2 file.")
F_EpRz = myfile["F_EpRz"]
else
error("Fast-ion distribution .jld2 file did not have any expected keys for the distribution (F_ps, f or F_EpRz).")
end