Integrate particle source

.github/workflows/python-publish.yml

@@ -20,8 +20,8 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         pip install setuptools wheel numpy
-    - name: Build the package
-      run: python setup.py sdist
+   # - name: Build the package
+    #  run: python setup.py sdist
     - name: Publish distribution 📦 to Test PyPI
       if: startsWith(github.ref, 'refs/tags')
       uses: pypa/gh-action-pypi-publish@master

aurora/adas_files.py

@@ -27,9 +27,13 @@
 import requests
 
 # location of the "adas_data" directory relative to this script:
-adas_data_dir = (
-    os.path.dirname(os.path.realpath(__file__)) + os.sep + "adas_data" + os.sep
-)
+if  "AURORA_ADAS_DIR" in os.environ:
+    adas_data_dir = os.environ['AURORA_ADAS_DIR']
+else:
+    adas_data_dir = (
+        os.path.dirname(os.path.realpath(__file__)) + os.sep + "adas_data" + os.sep
+    )
+
 
 
 def get_adas_file_loc(filename, filetype="adf11"):
@@ -75,12 +79,12 @@ def fetch_file(filename, filetype, loc):
 
         url = "https://open.adas.ac.uk/download/" + filetype + '/'
         if filetype == "adf11":
-            filename_mod = filename.split("_")[0] + os.sep + filename
+            filename_mod = filename.split("_")[0] + '/' + filename
 
         elif filetype in ["adf12", "adf13", "adf21", "adf22"]:
             filename_mod = (
                 filename.replace("#", "][").split("_")[0]
-                + os.sep
+                + '/'
                 + filename.replace("#", "][")
             )
 
@@ -99,15 +103,15 @@ def fetch_file(filename, filetype, loc):
             else:
                 # patterns may be different, attempt simple guess:
                 filename_mod = (
-                    filename.split("_")[0] + os.sep + filename.replace("#", "][")
+                    filename.split("_")[0] + '/' + filename.replace("#", "][")
                 )
         else:
             raise ValueError(
                 "ADAS file type/format not recognized.\n"
                 + "Could not find it or download it automatically!"
             )
 
-        r = requests.get(url + os.sep + filename_mod)
+        r = requests.get(url + '/' + filename_mod)
 
 
         if len(r.text) < 1000:
@@ -180,6 +184,7 @@ def adas_files_dict():
     files["H"]["brs"] = "brs05360.dat"
     files["H"]["pbs"] = "pbsx7_h.dat"
     files["H"]["prc"] = "prc89_h.dat"
+    files["D"] = files["H"]
     files["He"] = {}  # 2
     files["He"]["acd"] = "acd96_he.dat"
     files["He"]["scd"] = "scd96_he.dat"
@@ -382,6 +387,11 @@ def adas_files_dict():
     files["Ca"]["brs"] = "brs05360.dat"
     files["Ca"]["pbs"] = ""
     files["Ca"]["prc"] = ""
+    files["Ti"] = {}  # 22
+    files["Ti"]["acd"] = "acd00_ti.dat"
+    files["Ti"]["scd"] = "scd00_ti.dat"
+    files["Ti"]['prb'] = "prb00_ti.dat"
+    files["Ti"]['plt'] = "plt00_ti.dat"
     files["Fe"] = {}  # 26
     files["Fe"]["acd"] = "acd89_fe.dat"
     files["Fe"]["scd"] = "scd89_fe.dat"

aurora/atomic.py

@@ -143,7 +143,7 @@ def load(self):
                     self.MPRT.append(int(iprt.split("=")[1]))
                 except:
                     # some old files have different header
-                    self.Z.append(ind + 1)
+                    self.Z.append(ind)
                     self.MGRD.append(1)
                     self.MPRT.append(1)
 

aurora/coords.py

@@ -245,7 +245,7 @@ def rad_coord_transform(x, name_in, name_out, geqdsk):
     # Interpolate to transform coordiantes
     if name_in == "rhon":
         coord_in = rhon_ref
-    if name_in == "psin":
+    elif name_in == "psin":
         coord_in = psin_ref
     elif name_in == "rhop":
         coord_in = rhop_ref

aurora/core.py

@@ -1,5 +1,5 @@
 """This module includes the core class to set up simulations with :py:mod:`aurora`. The :py:class:`~aurora.core.aurora_sim` takes as input a namelist dictionary and a g-file dictionary (and possibly other optional argument) and allows creation of grids, interpolation of atomic rates and other steps before running the forward model.
-"""
+""" 
 # MIT License
 #
 # Copyright (c) 2021 Francesco Sciortino
@@ -66,7 +66,7 @@ def __init__(self, namelist, geqdsk=None):
         # make sure that any changes in namelist will not propagate back to the calling function
         self.namelist = deepcopy(namelist)
         self.geqdsk = geqdsk  # if None, minor (rvol_lcfs) and major radius (Raxis_cm) must be in namelist
-        self.kin_profs = self.namelist["kin_profs"]
+
         self.imp = namelist["imp"]
 
         # import here to avoid issues when building docs or package
@@ -107,7 +107,10 @@ def __init__(self, namelist, geqdsk=None):
             atom_files["ccd"] = self.namelist.get(
                 "ccd", adas_files.adas_files_dict()[self.imp]["ccd"]
             )
-
+            if atom_files["ccd"] is None:
+                raise Exception('Missing CCD ADF11 file!')
+
+
         if self.namelist.get("metastable_flag", False):
             atom_files["qcd"] = self.namelist.get(
                 "qcd", adas_files.adas_files_dict()[self.imp].get("qcd", None)
@@ -117,16 +120,20 @@ def __init__(self, namelist, geqdsk=None):
             )
 
         # now load ionization and recombination rates
+
         self.atom_data = atomic.get_atom_data(self.imp, files=atom_files)
 
         # allow for ion superstaging
         self.superstages = self.namelist.get("superstages", [])
 
         # set up radial and temporal grids
         self.setup_grids()
-
-        # set up kinetic profiles and atomic rates
-        self.setup_kin_profs_depts()
+
+        #interpolate kinetics profiles of the profiles are in namelist
+        if "kin_profs" in self.namelist:
+            self.kin_profs = self.namelist["kin_profs"]
+            # set up kinetic profiles and atomic rates
+            self.setup_kin_profs_depts()
 
     def reload_namelist(self, namelist=None):
         """(Re-)load namelist to update scalar variables."""
@@ -248,26 +255,23 @@ def setup_kin_profs_depts(self):
         if len(save_time) == 1:  # if time averaged profiles were used
             Sne0 = Sne0[:, [0]]  # 0th charge state (neutral)
 
+
+        # get time history and radial profiles separately
+        source_time_history = source_utils.get_source_time_history(
+            self.namelist, self.Raxis_cm, self.time_grid
+        )  # units of particles/s/cm for 1D source or particles/s/cm^3  for 2D source
+
+
         if self.namelist["source_type"] == "arbitrary_2d_source":
             # interpolate explicit source values on time and rhop grids of simulation
             # NB: explicit_source_vals should be in units of particles/s/cm^3 <-- ionization rate
             srho = self.namelist["explicit_source_rhop"]
-            stime = self.namelist["explicit_source_time"]
-            source = np.array(self.namelist["explicit_source_vals"]).T
-
-            spl = RectBivariateSpline(srho, stime, source, kx=1, ky=1)
-            # extrapolate by the nearest values
-            self.source_rad_prof = spl(
-                np.clip(self.rhop_grid, min(srho), max(srho)),
-                np.clip(self.time_grid, min(stime), max(stime)),
-            )
+            self.source_rad_prof  = interp1d(srho, source_time_history.T, 
+                                            axis=0,bounds_error=False,
+                                             fill_value=0)(self.rhop_grid)
             # Change units to particles/cm^3
             self.src_core = self.source_rad_prof / Sne0
         else:
-            # get time history and radial profiles separately
-            source_time_history = source_utils.get_source_time_history(
-                self.namelist, self.Raxis_cm, self.time_grid
-            )  # units of particles/s/cm
 
             # get radial profile of source function for each time step
             # dimensionless, normalized such that pnorm=1
@@ -284,7 +288,6 @@ def setup_kin_profs_depts(self):
             self.src_core = self.source_rad_prof * source_time_history[None, :]
 
         self.src_core = np.asfortranarray(self.src_core)
-
         # if wall_recycling>=0, return flows from the divertor are enabled
         if (
             self.wall_recycling >= 0
@@ -479,10 +482,19 @@ def set_time_dept_atomic_rates(self, superstages=[], metastables=False):
 
         # cache radiative & dielectronic recomb:
         self.alpha_RDR_rates = Rne
-
+      
         if self.namelist["cxr_flag"]:
             # Get an effective recombination rate by summing radiative & CX recombination rates
-            self.alpha_CX_rates = out.pop(0) * (self._n0 / self._ne)[:, None]
+            alpha_CX_rates = out.pop(0) * (self._n0 / self._ne)[:, None]
+            #metastable resolved CCD file exist only for C and N, use unresolved file for others
+            if metastables and alpha_CX_rates.shape[1] != Rne.shape[1]:
+                self.alpha_CX_rates = np.zeros_like(Rne)
+                for i, m in enumerate(self.atom_data["ccd"].meta_ind[1:]):
+                    j = self.atom_data["acd"].meta_ind.index(m)
+                    self.alpha_CX_rates[:,j] = alpha_CX_rates[:,i]
+            else:
+                self.alpha_CX_rates = alpha_CX_rates
+
             Rne = (
                 Rne + self.alpha_CX_rates
             )  # inplace addition would change also self.alpha_RDR_rates
@@ -491,13 +503,17 @@ def set_time_dept_atomic_rates(self, superstages=[], metastables=False):
             # include charge exchange between NBI neutrals and impurities
             self.nbi_cxr = interp1d(
                 self.namelist["nbi_cxr"]["rhop"],
-                self.namelist["nbi_cxr"]["vals"],
-                axis=0,
+                np.atleast_3d(self.namelist["nbi_cxr"]["vals"]).T,
                 bounds_error=False,
                 fill_value=0.0,
             )(self.rhop_grid)
+
+            if self.nbi_cxr.shape[2] > 1:
+                #time-dependent, times are switch times of beams, values are mean in between switch times
+                it = self.namelist["nbi_cxr"]["times"].searchsorted(self.time_grid)
+                self.nbi_cxr = self.nbi_cxr[it]
 
-            Rne = Rne + self.nbi_cxr.T[None, :, :]
+            Rne = Rne + self.nbi_cxr
 
         if len(superstages):
             self.superstages, Rne, Sne, self.fz_upstage = atomic.superstage_rates(

aurora/grids_utils.py

@@ -732,14 +732,15 @@ def estimate_boundary_distance(shot, device, time_ms):
             shot=shot,
             TDI="\\EFIT01::TOP.RESULTS.A_EQDSK.ORIGHT",
         )  # CMOD format, take ORIGHT
+        assert tmp.check()
     except Exception:
         tmp = OMFITmdsValue(
             server=device,
             treename="EFIT01",
             shot=shot,
             TDI="\\EFIT01::TOP.RESULTS.AEQDSK.GAPOUT",
         )  # useful variable for many other devices
-
+  
     time_vec = tmp.dim_of(0)
     data_r = tmp.data()
 
@@ -749,8 +750,7 @@ def estimate_boundary_distance(shot, device, time_ms):
 
     # take separation to limiter to be 2/3 of the separation to the wall boundary
     lim_sep = round(bound_sep * 2.0 / 3.0, 3)
-
-    return bound_sep, lim_sep
+    return bound_sep*100, lim_sep*100
 
 
 def vol_int(var, rvol_grid, pro_grid, Raxis_cm, rvol_max=None):

aurora/interp.py

@@ -109,10 +109,11 @@ def interp_quad(x, y, d, rLCFS, r):
     """Function 'interp' used for kinetic profiles."""
     f = interp1d(x, np.log(y), kind="quadratic", assume_sorted=True, copy=False)
     idx = np.searchsorted(r, rLCFS)
-    core = np.exp(f(np.clip(r[:idx] / rLCFS, 0, x[-1])))
+    core = np.exp(f(np.clip(r[:idx] / rLCFS, x[0], x[-1])))
     edge = core[..., [idx - 1]] * np.exp(
         -np.outer(1.0 / np.asarray(d), r[idx:] - r[idx - 1])
     )
+
 
     return np.concatenate([core, edge], axis=-1)
 
@@ -122,12 +123,12 @@ def interpa_quad(x, y, rLCFS, r):
     f = interp1d(
         x,
         np.log(y),
-        bounds_error=False,
         kind="quadratic",
         assume_sorted=True,
         copy=False,
     )
-    return np.exp(f(np.minimum(r / rLCFS, x[-1])))
+
+    return np.exp(f(np.clip(r / rLCFS, x[0], x[-1])))
 
 
 def interp(x, y, rLCFS, r):

aurora/radiation.py

@@ -670,7 +670,12 @@ def read_adf15(path, order=1):
     # Get the expected number of lines by reading the header:
     num_lines = int(header[0])
     spec = header[1].strip("/")
-    Z = int(''.join(header[2:-3]).strip(':'))
+
+    # header formating is not strictly defined and this can sometimes fail
+    try:
+        Z = int(''.join(header[2:-3]).strip(':'))
+    except:
+        Z = 0
 
     log10pec_dict = {}
     for i in range(num_lines):