Merge pull request #11 from ProjectTorreyPines/edge_prof

Finish overall preparation workflow structure and add utilities/setup to prepare for edge profiles

.JuliaFormatter.toml

@@ -0,0 +1,27 @@
+align_assignment = true
+align_conditional = true
+align_matrix = true
+align_pair_arrow = true
+align_struct_field = true
+always_for_in = true
+always_use_return = true
+annotate_untyped_fields_with_any = false
+conditional_to_if = true
+for_in_replacement = "∈"
+format_docstrings = true
+import_to_using = true
+indent = 4
+indent_submodule = true
+join_lines_based_on_source = true
+long_to_short_function_def = true
+margin = 88
+normalize_line_endings = "unix"
+pipe_to_function_call = true
+remove_extra_newlines = true
+separate_kwargs_with_semicolon = true
+surround_whereop_typeparameters = true
+trailing_comma = true
+whitespace_in_kwargs = false
+whitespace_ops_in_indices = false
+whitespace_typedefs = true
+yas_style_nesting = true

.github/workflows/format_check.yml

@@ -0,0 +1,36 @@
+name: Format Check
+
+on:
+  push:
+    branches: ["master", "dev", "format"]
+  pull_request:
+    branches: ["master", "dev"]
+jobs:
+  check:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        julia-version: [1.9.3]
+        julia-arch: [x86]
+        os: [ubuntu-latest]
+    steps:
+      - uses: julia-actions/setup-julia@latest
+        with:
+          version: ${{ matrix.julia-version }}
+
+      - uses: actions/checkout@v1
+      - name: Install JuliaFormatter and format
+        run: |
+          julia  -e 'using Pkg; Pkg.add(PackageSpec(name="JuliaFormatter"))'
+          julia  -e 'using JuliaFormatter; format(".", verbose=true)'
+      - name: Format check
+        run: |
+          julia -e '
+          out = Cmd(`git diff --name-only`) |> read |> String
+          if out == ""
+              exit(0)
+          else
+              @error "Some files have not been formatted !!!"
+              write(stdout, out)
+              exit(1)
+          end'

.gitignore

@@ -3,3 +3,4 @@
 # committed for packages, but should be committed for applications that require a static
 # environment.
 Manifest.toml
+sd_input_data.json

sample/ITER_Lore_2296_00000.dvc

@@ -1,10 +1,10 @@
-md5: fd0856e1e06a7ddab0b8e4aaddc170de
+md5: 377b6f3a1439c78e6a310739b4ca69d0
 frozen: true
 deps:
 - path: ITER_Lore_2296_00000
   repo:
-    url: [email protected]:ProjectTorreyPines/TestSamples.git
-    rev_lock: f905bd19e48bd54bc53416de9bbdaadda6654b9a
+    url: [email protected]:ProjectTorreyPines/SOLPSTestSamples.git
+    rev_lock: 35fd24298a4118eec6491dcd17d4e0ec97823e83
 outs:
 - md5: 48af1b920caaa6bc1788af8a36512638.dir
   size: 322802293

src/SD4SOLPS.jl

@@ -1,9 +1,9 @@
 module SD4SOLPS
 
-import OMAS
-import SOLPS2IMAS
-import EFIT
-import Interpolations
+using OMAS: OMAS
+using SOLPS2IMAS: SOLPS2IMAS
+using EFIT: EFIT
+using Interpolations: Interpolations
 #import GGDUtils
 
 export find_files_in_allowed_folders
@@ -21,22 +21,32 @@ Searches a list of allowed folders for a set of filenames that will provide info
 about the SOLPS case. Returns a list of filenames with complete paths.
 
 Example:
-SD4SOLPS.find_files_in_allowed_folders("<your samples folder>/D3D_Ma_184833_03600", eqdsk_file="g184833.03600")
+SD4SOLPS.find_files_in_allowed_folders(
+"<your samples folder>/D3D_Ma_184833_03600",
+eqdsk_file="g184833.03600",
+)
 """
-function find_files_in_allowed_folders(input_dirs...; eqdsk_file, recursive=true, allow_reduced_versions=false)
+function find_files_in_allowed_folders(
+    input_dirs...;
+    eqdsk_file,
+    recursive=true,
+    allow_reduced_versions=false,
+)
     files = ["b2fgmtry", "b2time.nc", "b2mn.dat", "gridspacedesc.yml", eqdsk_file]
-    reduced_files = ["b2fgmtry_red", "b2time_red.nc", "b2mn.dat", "gridspacedesc.yml", eqdsk_file]
+    reduced_files =
+        ["b2fgmtry_red", "b2time_red.nc", "b2mn.dat", "gridspacedesc.yml", eqdsk_file]
     output_files = fill("", length(files))
     if recursive
         dirs = []
-        for dir in input_dirs
-            dirs = append!(dirs, [subdir[1] for subdir in [item for item in walkdir(dir)]])
+        for dir ∈ input_dirs
+            dirs =
+                append!(dirs, [subdir[1] for subdir ∈ [item for item ∈ walkdir(dir)]])
         end
     else
         dirs = input_dirs
     end
-    for i in 1:length(files)
-        for dir in dirs
+    for i ∈ 1:length(files)
+        for dir ∈ dirs
             file = dir * "/" * files[i]
             reduced_file = dir * "/" * reduced_files[i]
             if isfile(file)
@@ -87,7 +97,8 @@ function geqdsk_to_imas(eqdsk_file, dd; time_index=1)
     p1.f_df_dpsi = g.ffprim
     p1.dpressure_dpsi = g.pprime
     p1.q = g.qpsi
-    if hasproperty(g, :rhovn)  # rhovn is not in the original EFIT.jl but is added on a branch
+    if hasproperty(g, :rhovn)
+        # rhovn is not in the original EFIT.jl but is added on a branch
         p1.rho_tor_norm = g.rhovn
     end
 
@@ -122,8 +133,7 @@ function geqdsk_to_imas(eqdsk_file, dd; time_index=1)
     limiter.type.description = "first wall"
     resize!(limiter.unit, 1)
     limiter.unit[1].outline.r = g.rlim
-    limiter.unit[1].outline.z = g.zlim
-
+    return limiter.unit[1].outline.z = g.zlim
 end
 
 """
@@ -136,14 +146,14 @@ r: R locations of desired output points / m
 z: Z locations of desired output points / m
 """
 function core_profile_2d(dd, prof_time_idx, eq_time_idx, quantity, r, z)
-    if !check_rho_1d(dd, time_slice=eq_time_idx)
+    if !check_rho_1d(dd; time_slice=eq_time_idx)
         throw(ArgumentError("Equilibrium rho profile in input DD was missing."))
     end
     prof = dd.core_profiles.profiles_1d[prof_time_idx]
     rho_prof = prof.grid.rho_tor_norm
     quantity_fields = split(quantity, ".")
     p = prof
-    for field in quantity_fields
+    for field ∈ quantity_fields
         p = getproperty(p, Symbol(field))
     end
     eqt = dd.equilibrium.time_slice[eq_time_idx]
@@ -182,14 +192,16 @@ function core_profile_2d(dd, prof_time_idx, eq_time_idx, quantity, r, z)
 
     # EQUILBRIUM (the connection from rho to R,Z via psi):
     # psin_eq_ext : 1D array of psi_N values that corresponds to rhon_eq_ext
-    # rhon_eq_ext : 1D array of rho_N values that corresponds to psin_eq_ext --> connects rho to everything else via psi
+    # rhon_eq_ext : 1D array of rho_N values that corresponds to psin_eq_ext
+    #       --> connects rho to everything else via psi
     # psinrz : 2D array of psi_N values that corresponds to r_eq and z_eq
     # r_eq and z_eq : 1D coordinate arrays that correspond to psinrz
 
     # OUTPUT INSTRUCTIONS:
     # r and z : coordinates of output points where values of p are desired
 
-    psi_at_requested_points = Interpolations.LinearInterpolation((r_eq, z_eq), psinrz).(r, z)
+    psi_at_requested_points =
+        Interpolations.LinearInterpolation((r_eq, z_eq), psinrz).(r, z)
     rhonpsi = Interpolations.LinearInterpolation(psin_eq_ext, rhon_eq_ext)
     rho_at_requested_points = rhonpsi.(psi_at_requested_points)
     itp = Interpolations.LinearInterpolation(rho_prof, p)
@@ -203,19 +215,41 @@ end
 Gathers SOLPS and EFIT files and loads them into IMAS structure. Extrapolates
 profiles as needed to get a complete picture.
 """
-function preparation(eqdsk_file, dirs...)
-    b2fgmtry, b2time, b2mn, gridspec, eqdsk = find_files_in_allowed_folders(
-        dirs..., eqdsk_file=eqdsk_file
-    )
+function preparation(
+    eqdsk_file,
+    dirs...;
+    core_method::String="simple",
+    edge_method::String="simple",
+    filename::String="sd_input_data",
+    output_format::String="json",
+)
+    b2fgmtry, b2time, b2mn, gridspec, eqdsk =
+        find_files_in_allowed_folders(dirs...; eqdsk_file=eqdsk_file)
     println("Found source files:")
     println("    b2fgmtry = ", b2fgmtry)
     println("    b2time = ", b2time)
     println("    b2mn.dat = ", b2mn)
     println("    gridspec = ", gridspec)
     println("    eqdsk = ", eqdsk)
+
     dd = SOLPS2IMAS.solps2imas(b2fgmtry, b2time, gridspec, b2mn)
     geqdsk_to_imas(eqdsk, dd)
-    println("Loaded data into IMAS DD")
+    println("Loaded input data into IMAS DD")
+
+    fill_in_extrapolated_core_profile!(dd, "electrons.density"; method=core_method)
+    fill_in_extrapolated_core_profile!(dd, "electrons.temperature"; method=core_method)
+    # ... more profiles here as they become available in b2time
+    println("Extrapolated core profiles")
+
+    fill_in_extrapolated_edge_profile!(dd, "electrons.density"; method=core_method)
+    # ... more profiles here
+    println("Extrapolated edge profiles (but not really (placeholder only))")
+
+    if output_format == "json"
+        OMAS.imas2json(dd, filename * ".json")
+    else
+        throw(ArgumentError(string("Unrecognized output format: ", output_format)))
+    end
     return dd
 end
 

src/repair_eq.jl

@@ -6,8 +6,8 @@ use this tool mainly for test cases, as a utility. For a real equilibrium,
 problems should be fixed properly.
 """
 
-import Contour
-import Statistics
+using Contour: Contour
+using Statistics: Statistics
 
 export add_rho_to_equilibrium!
 export check_rho_1d
@@ -40,7 +40,7 @@ function add_rho_to_equilibrium!(dd::OMAS.dd)
         println("No equilibrium time slices to work with; can't add rho")
         return
     end
-    for it = 1:nt
+    for it ∈ 1:nt
         eqt = dd.equilibrium.time_slice[it]
         b0 = dd.equilibrium.vacuum_toroidal_field.b0[it]
         r0 = dd.equilibrium.vacuum_toroidal_field.r0[it]
@@ -60,7 +60,9 @@ function add_rho_to_equilibrium!(dd::OMAS.dd)
             zeq = collect(eqt.profiles_2d[1].grid.dim2)
             if (length(req), length(zeq)) == size(psi2')
                 psi2 = Matrix(psi2')
-                println("transposed psi to make it compatible with r,z prior to contouring")
+                println(
+                    "transposed psi to make it compatible with r,z prior to contouring",
+                )
             end
             if (length(req), length(zeq)) != size(psi2)
                 println("Invalid equilibrium data. rho cannot be added.")
@@ -70,25 +72,38 @@ function add_rho_to_equilibrium!(dd::OMAS.dd)
                 println("    psi2d   : ", size(psi2))
                 return
             else
-                println("Eq looks okay ", (length(req), length(zeq)), ", ", size(psi2), ". ", (size(req), size(zeq)))
+                println(
+                    "Eq looks okay ",
+                    (length(req), length(zeq)),
+                    ", ",
+                    size(psi2),
+                    ". ",
+                    (size(req), size(zeq)),
+                )
             end
-            for j = 1:n
+            for j ∈ 1:n
                 contour_level = psi[j]
                 if j == n
-                    # The last contour has a X point if everything is lined up right, and that could get
-                    # weird. We want a contour that only takes the core boundary part of the separatrix.
+                    # The last contour has a X point if everything is lined up right,
+                    # and that could get weird. We want a contour that only takes the
+                    # core boundary part of the separatrix.
                     r = eqt.boundary.outline.r
                     z = eqt.boundary.outline.z
                 else
                     c = Contour.contour(req, zeq, psi2, contour_level)
                     clines = Contour.lines(c)
-                    line_avg_height = [Statistics.mean([clines[i].vertices[v][2] for v in 1:length(clines[i].vertices)]) for i in 1:length(clines)]
+                    line_avg_height = [
+                        Statistics.mean([
+                            clines[i].vertices[v][2] for
+                            v ∈ 1:length(clines[i].vertices)
+                        ]) for i ∈ 1:length(clines)
+                    ]
                     cl = clines[argmin(abs.(line_avg_height))]
                     # Now just do the 2D integral of B within the area of the contour
-                    r = [cl.vertices[v][1] for v in 1:length(cl.vertices)]
-                    z = [cl.vertices[v][2] for v in 1:length(cl.vertices)]
+                    r = [cl.vertices[v][1] for v ∈ 1:length(cl.vertices)]
+                    z = [cl.vertices[v][2] for v ∈ 1:length(cl.vertices)]
                 end
-                # Get the upper and lower halves of the path, each going from rmin to rmax.
+                # Get the upper & lower halves of the path, each going from rmin to rmax
                 rmin = minimum(r)
                 rmax = maximum(r)
                 irmin = argmin(r)
@@ -123,13 +138,17 @@ function add_rho_to_equilibrium!(dd::OMAS.dd)
                 z2i = Interpolations.LinearInterpolation(r2, z2)
                 rr = LinRange(rmin, rmax, 101)
                 rc = (rr[1:end-1] + rr[2:end]) / 2.0
-                integral_part_ = [log(rr[i+1]/rr[i]) * abs(z1i(rc[i]) - z2i(rc[i])) for i in 1:length(rc)]
+                integral_part_ = [
+                    log(rr[i+1] / rr[i]) * abs(z1i(rc[i]) - z2i(rc[i])) for
+                    i ∈ 1:length(rc)
+                ]
                 phi_ = r0 * b0 .* integral_part_
                 eqt.profiles_1d.phi[j] = sum(phi_)
             end
         end
         eqt.profiles_1d.rho_tor = sqrt.(eqt.profiles_1d.phi / (π * b0))
-        eqt.profiles_1d.rho_tor_norm = eqt.profiles_1d.rho_tor / eqt.profiles_1d.rho_tor[end]
+        eqt.profiles_1d.rho_tor_norm =
+            eqt.profiles_1d.rho_tor / eqt.profiles_1d.rho_tor[end]
     end
-    println("Rho has been added to the equilibrium")
+    return println("Rho has been added to the equilibrium")
 end