From 01cb714c85e75bbca279ee4209cf78c45dab9340 Mon Sep 17 00:00:00 2001 From: Dominic Power Date: Mon, 5 May 2025 18:11:49 -0400 Subject: [PATCH 01/10] Fixed plotting issue and added updown symmetry to SNL configs --- INGRID/gui/ingrid_gui.py | 2 + INGRID/ingrid.py | 26 +- INGRID/interpol.py | 7 +- INGRID/utils.py | 1438 +++++++++++++++++++++++++------------- 4 files changed, 990 insertions(+), 483 deletions(-) diff --git a/INGRID/gui/ingrid_gui.py b/INGRID/gui/ingrid_gui.py index af2813d..5e68d4d 100644 --- a/INGRID/gui/ingrid_gui.py +++ b/INGRID/gui/ingrid_gui.py @@ -472,6 +472,8 @@ def CreatePatches(self): self.AnalyzeTopology() self.controller.IngridSession.ConstructPatches() + self.controller.IngridSession.ApplyUpDownSymmetry() + self.controller.IngridSession.PlotPatches() def CreateSubgrid(self): diff --git a/INGRID/ingrid.py b/INGRID/ingrid.py index 034655e..686eb60 100644 --- a/INGRID/ingrid.py +++ b/INGRID/ingrid.py @@ -833,18 +833,18 @@ def PlotPsiNormBounds(self) -> None: nxpt = self.settings['grid_settings']['num_xpt'] if nxpt == 1: - Dic = {'psi_1': 'lime', + self.contour_colour_dic = {'psi_1': 'lime', 'psi_core': 'cyan', 'psi_pf_1': 'white'} elif nxpt == 2: - Dic = {'psi_core': 'cyan', + self.contour_colour_dic = {'psi_core': 'cyan', 'psi_1': 'lime', 'psi_2': 'fuchsia', 'psi_pf_1': 'white', 'psi_pf_2': 'yellow'} - for k, c in Dic.items(): - self.PsiNorm.PlotLevel(self.settings['grid_settings'][k], color=Dic[k], label=k) + for k, c in self.contour_colour_dic.items(): + self.PsiNorm.PlotLevel(self.settings['grid_settings'][k], color=self.contour_colour_dic[k], label=k) self.PsiNorm.PlotLevel(1.0, color='red', label='Primary Separatrix') if nxpt == 2: @@ -852,6 +852,12 @@ def PlotPsiNormBounds(self) -> None: self.PsiNorm.get_psi(self.xpt2[0], self.xpt2[1]), color='blue', label='Secondary Separatrix') handles, labels = self.PsiNorm.ax.get_legend_handles_labels() + contour_count = 0 + contour_colours = list(self.contour_colour_dic.values()) + ["red", "blue"] + for i in range(len(handles)): + if isinstance(handles[i], matplotlib.collections.PathCollection): + handles[i] = matplotlib.lines.Line2D([0],[0],color=contour_colours[contour_count]) + contour_count+=1 lookup = {label: handle for label, handle in zip(labels, handles)} try: self.PsiNorm.fig.legends[0].remove() @@ -859,7 +865,7 @@ def PlotPsiNormBounds(self) -> None: pass self.PsiNorm.fig.legend(handles=[handle for handle in lookup.values()], labels=[label for label in lookup.keys()], bbox_to_anchor=(0.5, 1), loc='upper center', - ncol=len([label for label in lookup.keys()]) // 3) + ncol=len([label for label in lookup.keys()]) // 3, facecolor="gray",framealpha=0.2) def PlotPsiNormMagReference(self, ax: object = None) -> None: """ @@ -1419,6 +1425,16 @@ def LoadPatches(self, fname: str = '') -> None: self.CurrentTopology.SetupPatchMatrix() self.CheckPatches() + def ApplyUpDownSymmetry(self) -> None: + if self.settings["grid_settings"]["up_down_symmetry"]: + if self.CurrentTopology.config == "LSN": + del self.CurrentTopology.patches["C1"] + del self.CurrentTopology.patches["C2"] + del self.CurrentTopology.patches["D1"] + del self.CurrentTopology.patches["D2"] + else: + raise Exception("Up/down symmetry not yet implemented for config " + self.CurrentTopology.config) + @_timer def ConstructGrid(self, NewFig: bool = True, ShowVertices: bool = False) -> None: """ diff --git a/INGRID/interpol.py b/INGRID/interpol.py index 6b324bd..75a8dc9 100644 --- a/INGRID/interpol.py +++ b/INGRID/interpol.py @@ -233,11 +233,10 @@ def PlotLevel(self: object, level: float = 1.0, color: str = 'red', label: str = indexing='ij') try: self.psi_levels[label].collections[0].remove() - self.psi_levels[label] = plt.contour(rgrid, zgrid, data, [float(level)], colors=color, label=label, linestyles=linestyles) - self.psi_levels[label].collections[0].set_label(label) except: - self.psi_levels[label] = plt.contour(rgrid, zgrid, data, [float(level)], colors=color, label=label, linestyles=linestyles) - self.psi_levels[label].collections[0].set_label(label) + pass + self.psi_levels[label] = plt.contour(rgrid, zgrid, data, [float(level)], colors=color, label=label, linestyles=linestyles) + self.psi_levels[label].collections[0].set_label(label) def plot_data(self: object, nlevs: int = 30, interactive: bool = True, fig: object = None, ax: object = None, view_mode: str = 'filled', refined: bool = True, refine_factor: int = 10): diff --git a/INGRID/utils.py b/INGRID/utils.py index 69958dd..518df7d 100644 --- a/INGRID/utils.py +++ b/INGRID/utils.py @@ -10,6 +10,7 @@ from __future__ import print_function, division, absolute_import import numpy as np import matplotlib + try: matplotlib.use("TkAgg") except: @@ -26,7 +27,7 @@ from INGRID.geometry import Point, Line, Patch, orientation_between -class IngridUtils(): +class IngridUtils: """ The base class for :class:`ingrid.Ingrid` that handles backend management of key Ingrid capabilities. This class can be directly utilized by advanced users @@ -119,11 +120,11 @@ def __init__(self, settings: dict = {}, **kwargs): self.InputFile = None self.config = None self.settings_lookup = [ - 'dir_settings', - 'grid_settings', - 'integrator_settings', - 'target_plates', - 'DEBUG' + "dir_settings", + "grid_settings", + "integrator_settings", + "target_plates", + "DEBUG", ] self.SetDefaultSettings() self.PopulateSettings(settings) @@ -160,147 +161,125 @@ def SetDefaultSettings(self) -> None: """ self.default_grid_settings = { - 'num_xpt': 1, - 'nlevs': 30, - 'view_mode': 'filled', - 'psi_1': 0.0, - 'psi_core': 0.0, - 'psi_pf_1': 0.0, - 'psi_pf_2': 0.0, - 'psi_1': 0.0, - 'psi_2': 0.0, - 'rmagx': 0.0, - 'zmagx': 0.0, - 'rxpt': 0.0, - 'zxpt': 0.0, - 'rxpt2': 0.0, - 'zxpt2': 0.0, - 'guard_cell_eps': 1.0e-3, - 'grid_generation': { - 'np_default': 2, - 'nr_default': 2, - 'poloidal_f_default': 'x, x', - 'radial_f_default': 'x, x', - 'distortion_correction': { - 'all': { - 'theta_min': 80.0, - 'theta_max': 120.0, - 'resolution': 1000, - 'active': False + "num_xpt": 1, + "nlevs": 30, + "view_mode": "filled", + "psi_1": 0.0, + "psi_core": 0.0, + "psi_pf_1": 0.0, + "psi_pf_2": 0.0, + "psi_1": 0.0, + "psi_2": 0.0, + "rmagx": 0.0, + "zmagx": 0.0, + "rxpt": 0.0, + "zxpt": 0.0, + "rxpt2": 0.0, + "zxpt2": 0.0, + "guard_cell_eps": 1.0e-3, + "grid_generation": { + "np_default": 2, + "nr_default": 2, + "poloidal_f_default": "x, x", + "radial_f_default": "x, x", + "distortion_correction": { + "all": { + "theta_min": 80.0, + "theta_max": 120.0, + "resolution": 1000, + "active": False, }, }, }, - 'patch_generation': { - 'core_split_point_ratio': 0.5, - 'pf_split_point_ratio': 0.5, - 'strike_pt_loc': 'limiter', - 'rmagx_shift': 0.0, - 'zmagx_shift': 0.0, - 'magx_tilt_1': 0.0, - 'magx_tilt_2': 0.0, - 'use_xpt1_W': False, - 'use_xpt1_E': False, - 'use_xpt2_W': False, - 'use_xpt2_E': False, - 'xpt1_W_tilt': -0.785398, # All values of pi / 4 radians. - 'xpt1_E_tilt': 0.785398, - 'xpt2_W_tilt': -0.785398, - 'xpt2_E_tilt': 0.785398, - } + "patch_generation": { + "core_split_point_ratio": 0.5, + "pf_split_point_ratio": 0.5, + "strike_pt_loc": "limiter", + "rmagx_shift": 0.0, + "zmagx_shift": 0.0, + "magx_tilt_1": 0.0, + "magx_tilt_2": 0.0, + "use_xpt1_W": False, + "use_xpt1_E": False, + "use_xpt2_W": False, + "use_xpt2_E": False, + "xpt1_W_tilt": -0.785398, # All values of pi / 4 radians. + "xpt1_E_tilt": 0.785398, + "xpt2_W_tilt": -0.785398, + "xpt2_E_tilt": 0.785398, + }, + "up_down_symmetry": False } self.default_integrator_settings = { - 'dt': 0.01, - 'eps': 5e-5, - 'first_step': 1e-5, - 'step_ratio': 0.02, - 'tol': 5e-3, - 'max_step': 0.064 + "dt": 0.01, + "eps": 5e-5, + "first_step": 1e-5, + "step_ratio": 0.02, + "tol": 5e-3, + "max_step": 0.064, } self.default_target_plate_settings = { - 'plate_E1': { - 'file': '', - 'rshift': 0.0, - 'zshift': 0.0 - }, - - 'plate_E2': { - 'file': '', - 'rshift': 0.0, - 'zshift': 0.0 - }, - - 'plate_W1': { - 'file': '', - 'rshift': 0.0, - 'zshift': 0.0 - }, - - 'plate_W2': { - 'file': '', - 'rshift': 0.0, - 'zshift': 0.0 - }, + "plate_E1": {"file": "", "rshift": 0.0, "zshift": 0.0}, + "plate_E2": {"file": "", "rshift": 0.0, "zshift": 0.0}, + "plate_W1": {"file": "", "rshift": 0.0, "zshift": 0.0}, + "plate_W2": {"file": "", "rshift": 0.0, "zshift": 0.0}, } self.default_limiter_settings = { - 'file': '', - 'use_efit_bounds': False, - 'efit_buffer_r': 1e-2, - 'efit_buffer_z': 1e-2, - 'rshift': 0.0, - 'zshift': 0.0 + "file": "", + "use_efit_bounds": False, + "efit_buffer_r": 1e-2, + "efit_buffer_z": 1e-2, + "rshift": 0.0, + "zshift": 0.0, } self.default_patch_data_settings = { - 'file': '', - 'use_file': False, - 'preferences': { - 'new_file': False, - 'new_fname': '' - } + "file": "", + "use_file": False, + "preferences": {"new_file": False, "new_fname": ""}, } self.default_DEBUG_settings = { - 'visual': { - 'find_NSEW': False, - 'patch_map': False, - 'subgrid': False, - 'gridue': False, - 'SF_analysis': False + "visual": { + "find_NSEW": False, + "patch_map": False, + "subgrid": False, + "gridue": False, + "SF_analysis": False, + }, + "verbose": { + "patch_generation": False, + "grid_generation": False, + "SF_analysis": False, }, - - 'verbose': { - 'patch_generation': False, - 'grid_generation': False, - 'SF_analysis': False - } } self.default_dir_settings = { - 'eqdsk': '.', - 'limiter': '.', - 'patch_data': '.', - 'target_plates': '.' + "eqdsk": ".", + "limiter": ".", + "patch_data": ".", + "target_plates": ".", } self.default_values_lookup = { - 'eqdsk': '', - 'dir_settings': self.default_dir_settings, - 'grid_settings': self.default_grid_settings, - 'integrator_settings': self.default_integrator_settings, - 'target_plates': self.default_target_plate_settings, - 'limiter': self. default_limiter_settings, - 'patch_data': self.default_patch_data_settings, - 'DEBUG': self.default_DEBUG_settings + "eqdsk": "", + "dir_settings": self.default_dir_settings, + "grid_settings": self.default_grid_settings, + "integrator_settings": self.default_integrator_settings, + "target_plates": self.default_target_plate_settings, + "limiter": self.default_limiter_settings, + "patch_data": self.default_patch_data_settings, + "DEBUG": self.default_DEBUG_settings, } self.PlateData = { - 'plate_W1': {}, - 'plate_E1': {}, - 'plate_W2': {}, - 'plate_E2': {} + "plate_W1": {}, + "plate_E1": {}, + "plate_W2": {}, + "plate_E2": {}, } def ProcessKeywords(self, **kwargs) -> None: @@ -308,27 +287,27 @@ def ProcessKeywords(self, **kwargs) -> None: Process kwargs and set all file paths accordingly. """ for k, v in kwargs.items(): - if k == 'InputFile' or k == 'yaml': - print('# Processing Input File:', v) + if k == "InputFile" or k == "yaml": + print("# Processing Input File:", v) self.InputFile = v self.PopulateSettings(self.ReadYamlFile(v)) continue - if k == 'W1TargetFile' or k == 'w1': - self.settings['target_plates']['plate_W1']['file'] = v + if k == "W1TargetFile" or k == "w1": + self.settings["target_plates"]["plate_W1"]["file"] = v continue - if k == 'E1TargetFile' or k == 'e1': - self.settings['target_plates']['plate_E1']['file'] = v + if k == "E1TargetFile" or k == "e1": + self.settings["target_plates"]["plate_E1"]["file"] = v continue - if k == 'E2TargetFile' or k == 'e2': - self.settings['target_plates']['plate_E2']['file'] = v + if k == "E2TargetFile" or k == "e2": + self.settings["target_plates"]["plate_E2"]["file"] = v continue - if k == 'W2TargetFile' or k == 'w2': - self.settings['target_plates']['plate_W2']['file'] = v + if k == "W2TargetFile" or k == "w2": + self.settings["target_plates"]["plate_W2"]["file"] = v continue - if k in ['LimiterFile', 'limiter', 'wall']: - self.settings['limiter']['file'] = v - if k == 'EqFile' or k == 'eq': - self.settings['eqdsk'] = v + if k in ["LimiterFile", "limiter", "wall"]: + self.settings["limiter"]["file"] = v + if k == "EqFile" or k == "eq": + self.settings["eqdsk"] = v continue print('Keyword "' + k + '" unknown and ignored...') @@ -358,7 +337,9 @@ def _check_settings_input(input_settings, comparison): raise_assertion = True if raise_assertion is True: - raise ValueError(f'Invalid entries {items} in provided settings file f{self.InputFile}. Remove invalid entries listed.') + raise ValueError( + f"Invalid entries {items} in provided settings file f{self.InputFile}. Remove invalid entries listed." + ) def _get_default_values(item, sub_item=None, attribute=None): """ @@ -369,7 +350,11 @@ def _get_default_values(item, sub_item=None, attribute=None): try: default_values = self.default_values_lookup[item] except KeyError: - print(f'Key ''{item}'' not recognized... Add default values to source code for support.') + print( + f"Key " + "{item}" + " not recognized... Add default values to source code for support." + ) return None if item and sub_item and attribute: @@ -379,7 +364,9 @@ def _get_default_values(item, sub_item=None, attribute=None): elif item: return self.default_values_lookup[item] - _check_settings_input(input_settings=settings, comparison=self.default_values_lookup) + _check_settings_input( + input_settings=settings, comparison=self.default_values_lookup + ) # First level entries within YAML file dump. for item in self.default_values_lookup.keys(): @@ -391,13 +378,20 @@ def _get_default_values(item, sub_item=None, attribute=None): settings[item] = _get_default_values(item) if verbose: print('Could not find "{}" in YAML file.'.format(item)) - print('Populated "{}" with default value of "{}".\n'.format(item, settings[item])) + print( + 'Populated "{}" with default value of "{}".\n'.format( + item, settings[item] + ) + ) continue - if item in ['eqdsk']: + if item in ["eqdsk"]: continue - _check_settings_input(input_settings=settings[item], comparison=self.default_values_lookup[item]) + _check_settings_input( + input_settings=settings[item], + comparison=self.default_values_lookup[item], + ) # Second level entries within YAML file dump. @@ -407,25 +401,45 @@ def _get_default_values(item, sub_item=None, attribute=None): except KeyError: settings[item][sub_item] = _get_default_values(item, sub_item) if verbose: - print('Could not find "{}/{}" in YAML file.'.format(item, sub_item)) - print('Populated "{}/{}" with default value of "{}".\n'.format(item, sub_item, settings[item][sub_item])) + print( + 'Could not find "{}/{}" in YAML file.'.format( + item, sub_item + ) + ) + print( + 'Populated "{}/{}" with default value of "{}".\n'.format( + item, sub_item, settings[item][sub_item] + ) + ) continue - if item in ['grid_settings', 'target_plates'] \ - and sub_item in ['patch_generation', 'grid_generation', 'plate_E1', 'plate_W1', 'plate_E2', 'plate_W2']: - for plate_attribute in self.default_values_lookup[item][sub_item].keys(): + if item in ["grid_settings", "target_plates"] and sub_item in [ + "patch_generation", + "grid_generation", + "plate_E1", + "plate_W1", + "plate_E2", + "plate_W2", + ]: + for plate_attribute in self.default_values_lookup[item][ + sub_item + ].keys(): try: if settings[item][sub_item][plate_attribute] is None: - settings[item][sub_item][plate_attribute] = _get_default_values(item, sub_item, plate_attribute) + settings[item][sub_item][plate_attribute] = ( + _get_default_values(item, sub_item, plate_attribute) + ) except: - settings[item][sub_item][plate_attribute] = _get_default_values(item, sub_item, plate_attribute) + settings[item][sub_item][plate_attribute] = ( + _get_default_values(item, sub_item, plate_attribute) + ) # Update references to YAML entries. self.settings = settings - self.grid_settings = settings['grid_settings'] - self.integrator_settings = settings['integrator_settings'] - self.target_plates = settings['target_plates'] - self.DEBUG = settings['DEBUG'] + self.grid_settings = settings["grid_settings"] + self.integrator_settings = settings["integrator_settings"] + self.target_plates = settings["target_plates"] + self.DEBUG = settings["DEBUG"] self.ProcessPaths() def ProcessPaths(self) -> None: @@ -433,26 +447,55 @@ def ProcessPaths(self) -> None: Update settings by pre-pending path entries to all file entries. """ for k in self.default_dir_settings.keys(): - path_obj = Path(self.settings['dir_settings'][k]) + path_obj = Path(self.settings["dir_settings"][k]) if path_obj.is_dir() is False: - v = self.settings['dir_settings'][k] - raise ValueError(f"# Error processing directory provided for entry '{k}'. Entry '{v}' is not a valid directory.") - - if k == 'eqdsk': - self.settings['eqdsk'] = str((path_obj / self.settings['eqdsk']).absolute()) + v = self.settings["dir_settings"][k] + raise ValueError( + f"# Error processing directory provided for entry '{k}'. Entry '{v}' is not a valid directory." + ) + + if k == "eqdsk": + self.settings["eqdsk"] = str( + (path_obj / self.settings["eqdsk"]).absolute() + ) continue - if k == 'limiter': - self.settings['limiter']['file'] = str((path_obj / self.settings['limiter']['file']).absolute()) + if k == "limiter": + self.settings["limiter"]["file"] = str( + (path_obj / self.settings["limiter"]["file"]).absolute() + ) continue - if k == 'target_plates': - self.settings['target_plates']['plate_W1']['file'] = str((path_obj / self.settings['target_plates']['plate_W1']['file']).absolute()) - self.settings['target_plates']['plate_E1']['file'] = str((path_obj / self.settings['target_plates']['plate_E1']['file']).absolute()) - self.settings['target_plates']['plate_W2']['file'] = str((path_obj / self.settings['target_plates']['plate_W2']['file']).absolute()) - self.settings['target_plates']['plate_E2']['file'] = str((path_obj / self.settings['target_plates']['plate_E2']['file']).absolute()) + if k == "target_plates": + self.settings["target_plates"]["plate_W1"]["file"] = str( + ( + path_obj / self.settings["target_plates"]["plate_W1"]["file"] + ).absolute() + ) + self.settings["target_plates"]["plate_E1"]["file"] = str( + ( + path_obj / self.settings["target_plates"]["plate_E1"]["file"] + ).absolute() + ) + self.settings["target_plates"]["plate_W2"]["file"] = str( + ( + path_obj / self.settings["target_plates"]["plate_W2"]["file"] + ).absolute() + ) + self.settings["target_plates"]["plate_E2"]["file"] = str( + ( + path_obj / self.settings["target_plates"]["plate_E2"]["file"] + ).absolute() + ) continue - if k == 'patch_data': - self.settings['patch_data']['file'] = str((path_obj / self.settings['patch_data']['file']).absolute()) - self.settings['patch_data']['preferences']['new_fname'] = str((path_obj / self.settings['patch_data']['preferences']['new_fname']).absolute()) + if k == "patch_data": + self.settings["patch_data"]["file"] = str( + (path_obj / self.settings["patch_data"]["file"]).absolute() + ) + self.settings["patch_data"]["preferences"]["new_fname"] = str( + ( + path_obj + / self.settings["patch_data"]["preferences"]["new_fname"] + ).absolute() + ) continue def LoadGEQDSK(self, geqdsk_path: str) -> None: @@ -462,59 +505,66 @@ def LoadGEQDSK(self, geqdsk_path: str) -> None: Saves the boundary information and generates an EfitData instance. """ - with open(geqdsk_path, 'r') as f: + with open(geqdsk_path, "r") as f: geqdsk_data = geqdsk.read(f) # # Extract quantities needed to initialize EfitData class # - nx = geqdsk_data['nx'] - ny = geqdsk_data['ny'] - rdim = geqdsk_data['rdim'] - zdim = geqdsk_data['zdim'] - zmid = geqdsk_data['zmid'] - rleft = geqdsk_data['rleft'] - rcenter = geqdsk_data['rcentr'] - bcenter = geqdsk_data['bcentr'] - rmagx = geqdsk_data['rmagx'] - zmagx = geqdsk_data['zmagx'] - rlimiter = geqdsk_data['rlim'] - zlimiter = geqdsk_data['zlim'] - psi = geqdsk_data['psi'] + nx = geqdsk_data["nx"] + ny = geqdsk_data["ny"] + rdim = geqdsk_data["rdim"] + zdim = geqdsk_data["zdim"] + zmid = geqdsk_data["zmid"] + rleft = geqdsk_data["rleft"] + rcenter = geqdsk_data["rcentr"] + bcenter = geqdsk_data["bcentr"] + rmagx = geqdsk_data["rmagx"] + zmagx = geqdsk_data["zmagx"] + rlimiter = geqdsk_data["rlim"] + zlimiter = geqdsk_data["zlim"] + psi = geqdsk_data["psi"] # # Derived values # rmin = rleft rmax = rmin + rdim - zmin = (zmid - 0.5 * zdim) + zmin = zmid - 0.5 * zdim zmax = zmin + zdim # # Reproduce efit grid # self.PsiUNorm = EfitData( - rmin=rmin, - rmax=rmax, - nr=nx, - zmin=zmin, - zmax=zmax, - nz=ny, - rcenter=rcenter, - bcenter=bcenter, - rlimiter=rlimiter, - zlimiter=zlimiter, - rmagx=rmagx, - zmagx=zmagx, - name='Efit Data', - parent=self) - - self.PsiUNorm.init_bivariate_spline(self.PsiUNorm.r[:, 0], - self.PsiUNorm.z[0, :], - psi) - - if self.settings['grid_settings']['rmagx'] is None or self.settings['grid_settings']['zmagx'] is None: - self.settings['grid_settings']['rmagx'], self.settings['grid_settings']['zmagx'] = (self.PsiUNorm.rmagx, self.PsiUNorm.zmagx) + rmin=rmin, + rmax=rmax, + nr=nx, + zmin=zmin, + zmax=zmax, + nz=ny, + rcenter=rcenter, + bcenter=bcenter, + rlimiter=rlimiter, + zlimiter=zlimiter, + rmagx=rmagx, + zmagx=zmagx, + name="Efit Data", + parent=self, + ) + + self.PsiUNorm.init_bivariate_spline( + self.PsiUNorm.r[:, 0], self.PsiUNorm.z[0, :], psi + ) + + if ( + self.settings["grid_settings"]["rmagx"] is None + or self.settings["grid_settings"]["zmagx"] is None + ): + ( + self.settings["grid_settings"]["rmagx"], + self.settings["grid_settings"]["zmagx"], + ) = (self.PsiUNorm.rmagx, self.PsiUNorm.zmagx) # # Save the stored data in the object @@ -523,10 +573,11 @@ def LoadGEQDSK(self, geqdsk_path: str) -> None: # # Update the eqdsk file referenced in settings to that of the loaded data # - self.settings['eqdsk'] = geqdsk_path - + self.settings["eqdsk"] = geqdsk_path - def ParseTxtCoordinates(self, fpath: str, rshift: float = 0.0, zshift: float = 0.0) -> tuple: + def ParseTxtCoordinates( + self, fpath: str, rshift: float = 0.0, zshift: float = 0.0 + ) -> tuple: """ Extract the (R,Z) coordinates from a .txt file. @@ -578,37 +629,52 @@ def ParseTxtCoordinates(self, fpath: str, rshift: float = 0.0, zshift: float = 0 R, Z = [], [] - if Path(fpath).is_file() and Path(fpath).suffix in ['.txt']: + if Path(fpath).is_file() and Path(fpath).suffix in [".txt"]: try: with open(fpath) as f: for line in f: - if line.startswith('#'): + if line.startswith("#"): # move along to the next iteration # this simulates a comment continue # we deal with separator ',' or ' ' - point = line.replace('D', 'e').replace('(', '').replace(')', '').strip() - if point.count(',') > 0: - x = float(point.split(',')[0]) - y = float(point.split(',')[1]) + point = ( + line.replace("D", "e") + .replace("(", "") + .replace(")", "") + .strip() + ) + if point.count(",") > 0: + x = float(point.split(",")[0]) + y = float(point.split(",")[1]) else: - x = float(point.split(' ')[0]) - y = float(point.split(' ')[1]) + x = float(point.split(" ")[0]) + y = float(point.split(" ")[1]) R.append(x + rshift) Z.append(y + zshift) except: - raise IOError(f"# Error occured when reading data from file {fpath}:\t 'open(fpath)' error") + raise IOError( + f"# Error occured when reading data from file {fpath}:\t 'open(fpath)' error" + ) else: - raise ValueError(f"# Error occur when reading data from file {fpath}:\t file does not exist or is not of extension '*.txt'") + raise ValueError( + f"# Error occur when reading data from file {fpath}:\t file does not exist or is not of extension '*.txt'" + ) return R, Z - def SetLimiter(self, fpath: str = '', coordinates: list = [], rshift: float = 0.0, zshift: float = 0.0) -> None: + def SetLimiter( + self, + fpath: str = "", + coordinates: list = [], + rshift: float = 0.0, + zshift: float = 0.0, + ) -> None: """ Instantiate the class Line object that represents the tokamak limiter. @@ -635,28 +701,30 @@ def SetLimiter(self, fpath: str = '', coordinates: list = [], rshift: float = 0. """ if rshift is None: - rshift = self.settings['limiter']['rshift'] + rshift = self.settings["limiter"]["rshift"] if zshift is None: - zshift = self.settings['limiter']['zshift'] + zshift = self.settings["limiter"]["zshift"] - use_efit_bounds = self.settings['limiter']['use_efit_bounds'] + use_efit_bounds = self.settings["limiter"]["use_efit_bounds"] # # TODO: Refactor the codebase to leverage numpy arrays.... # coordinates = np.array(coordinates).T - + # # If we have no coordinates or fpath provided, we need to # leverage efit boundaries. Overrides any user settings. # - if coordinates.shape[0] == 0 and fpath == '': + if coordinates.shape[0] == 0 and fpath == "": use_efit_bounds = True - if fpath not in ['', '.']: + if fpath not in ["", "."]: try: - print('# Processing file for limiter data : {}'.format(fpath)) - self.geqdsk_data['rlim'], self.geqdsk_data['rlim'] = self.ParseTxtCoordinates(fpath) + print("# Processing file for limiter data : {}".format(fpath)) + self.geqdsk_data["rlim"], self.geqdsk_data["rlim"] = ( + self.ParseTxtCoordinates(fpath) + ) except: raise ValueError(f"# Error in method 'SetLimiter' with fpath={fpath}") @@ -664,16 +732,16 @@ def SetLimiter(self, fpath: str = '', coordinates: list = [], rshift: float = 0. # List of coordinates provided for initialization # elif len(coordinates) > 0: - self.geqdsk_data['rlim'] = coordinates[:, 0] + rshift - self.geqdsk_data['zlim'] = coordinates[:, 1] + zshift + self.geqdsk_data["rlim"] = coordinates[:, 0] + rshift + self.geqdsk_data["zlim"] = coordinates[:, 1] + zshift # # Empty list of coordinates falls back on using eqdsk limiter settings # else: - self.LoadGEQDSK(geqdsk_path=self.settings['eqdsk']) - self.geqdsk_data['rlim'] += rshift - self.geqdsk_data['zlim'] += zshift + self.LoadGEQDSK(geqdsk_path=self.settings["eqdsk"]) + self.geqdsk_data["rlim"] += rshift + self.geqdsk_data["zlim"] += zshift if use_efit_bounds: @@ -682,8 +750,8 @@ def SetLimiter(self, fpath: str = '', coordinates: list = [], rshift: float = 0. # Buffer values shrink the bounding box to allow us to avoid stepping # outside the domain during integration/line-tracing # - rbuff = self.settings['limiter']['efit_buffer_r'] - zbuff = self.settings['limiter']['efit_buffer_z'] + rbuff = self.settings["limiter"]["efit_buffer_r"] + zbuff = self.settings["limiter"]["efit_buffer_z"] # # Define the bounding box vertices as: @@ -700,13 +768,17 @@ def SetLimiter(self, fpath: str = '', coordinates: list = [], rshift: float = 0. # # Define the simple limiter path and translate by shift values # - limiter_path = np.vstack([LL, LR, UR, UL, LL]) - self.geqdsk_data['rlim'] = limiter_path[:, 0] + rshift - self.geqdsk_data['zlim'] = limiter_path[:, 1] + zshift + limiter_path = np.vstack([LL, LR, UR, UL, LL]) + self.geqdsk_data["rlim"] = limiter_path[:, 0] + rshift + self.geqdsk_data["zlim"] = limiter_path[:, 1] + zshift - self.LimiterData = Line([Point(p) for p in zip(self.geqdsk_data['rlim'], self.geqdsk_data['zlim'])]) + self.LimiterData = Line( + [Point(p) for p in zip(self.geqdsk_data["rlim"], self.geqdsk_data["zlim"])] + ) - def SetTargetPlate(self, settings: dict, rshift: float = 0.0, zshift: float = 0.0) -> None: + def SetTargetPlate( + self, settings: dict, rshift: float = 0.0, zshift: float = 0.0 + ) -> None: """ Initialize a target plate Line object. @@ -772,14 +844,14 @@ def SetTargetPlate(self, settings: dict, rshift: float = 0.0, zshift: float = 0. v = _v break - if k.lower() in ['w1', 'westplate1', 'plate_w1']: - k = 'plate_W1' - elif k.lower() in ['e1', 'eastplate1', 'plate_e1']: - k = 'plate_E1' - elif k.lower() in ['w2', 'westplate2', 'plate_w2']: - k = 'plate_W2' - elif k.lower() in ['e2', 'eastplate2', 'plate_e2']: - k = 'plate_E2' + if k.lower() in ["w1", "westplate1", "plate_w1"]: + k = "plate_W1" + elif k.lower() in ["e1", "eastplate1", "plate_e1"]: + k = "plate_E1" + elif k.lower() in ["w2", "westplate2", "plate_w2"]: + k = "plate_W2" + elif k.lower() in ["e2", "eastplate2", "plate_e2"]: + k = "plate_E2" else: raise ValueError(f"# Invalid key '{k}' provided for 'SetTargetPlate'") @@ -791,7 +863,9 @@ def SetTargetPlate(self, settings: dict, rshift: float = 0.0, zshift: float = 0. data = np.array([c for c in zip(R, Z)]) a, index = np.unique(data, return_index=True, axis=0) index.sort() - self.PlateData[k] = Line([Point(x + rshift, y + zshift) for x, y in data[index]]) + self.PlateData[k] = Line( + [Point(x + rshift, y + zshift) for x, y in data[index]] + ) def OrderTargetPlate(self, plate_key: str) -> None: """ @@ -823,7 +897,9 @@ def OrderTargetPlate(self, plate_key: str) -> None: k = plate_key if not isinstance(self.PlateData.get(k), Line): - raise ValueError(f"# Plate '{k}' is not loaded as a Line object. Make sure 'SetGeometry({{'{k}' : ... }})' has been called.") + raise ValueError( + f"# Plate '{k}' is not loaded as a Line object. Make sure 'SetGeometry({{'{k}' : ... }})' has been called." + ) plate = self.PlateData[k] @@ -854,8 +930,12 @@ def OrderTargetPlate(self, plate_key: str) -> None: # Check the angle between the plate endpoints with # tail fixed on the magnetic axis - v_reference = np.array([self.settings['grid_settings']['rmagx'], - self.settings['grid_settings']['zmagx']]) + v_reference = np.array( + [ + self.settings["grid_settings"]["rmagx"], + self.settings["grid_settings"]["zmagx"], + ] + ) v_start = np.array([start.x, start.y]) v_end = np.array([end.x, end.y]) @@ -936,8 +1016,12 @@ def OrderLimiter(self) -> None: # Check the angle between the plate endpoints with # tail fixed on the magnetic axis - v_reference = np.array([self.settings['grid_settings']['rmagx'], - self.settings['grid_settings']['zmagx']]) + v_reference = np.array( + [ + self.settings["grid_settings"]["rmagx"], + self.settings["grid_settings"]["zmagx"], + ] + ) v_start = np.array([start.x, start.y]) v_end = np.array([end.x, end.y]) @@ -961,8 +1045,8 @@ def FindMagAxis(self, r: float, z: float) -> None: Z coordinate of magnetic-axis guess. """ sol = root(self.PsiUNorm.Gradient, [r, z]) - self.settings['grid_settings']['rmagx'] = sol.x[0] - self.settings['grid_settings']['zmagx'] = sol.x[1] + self.settings["grid_settings"]["rmagx"] = sol.x[0] + self.settings["grid_settings"]["zmagx"] = sol.x[1] self.magx = (sol.x[0], sol.x[1]) def FindXPoint(self, r: float, z: float) -> None: @@ -978,8 +1062,8 @@ def FindXPoint(self, r: float, z: float) -> None: Z coordinate of primary x-point guess. """ sol = root(self.PsiUNorm.Gradient, [r, z]) - self.settings['grid_settings']['rxpt'] = sol.x[0] - self.settings['grid_settings']['zxpt'] = sol.x[1] + self.settings["grid_settings"]["rxpt"] = sol.x[0] + self.settings["grid_settings"]["zxpt"] = sol.x[1] self.xpt1 = (sol.x[0], sol.x[1]) def FindXPoint2(self, r: float, z: float) -> None: @@ -995,12 +1079,12 @@ def FindXPoint2(self, r: float, z: float) -> None: Z coordinate of secondary x-point guess. """ sol = root(self.PsiUNorm.Gradient, [r, z]) - self.settings['grid_settings']['rxpt2'] = sol.x[0] - self.settings['grid_settings']['zxpt2'] = sol.x[1] + self.settings["grid_settings"]["rxpt2"] = sol.x[0] + self.settings["grid_settings"]["zxpt2"] = sol.x[1] self.xpt2 = (sol.x[0], sol.x[1]) def _find_roots(self, tk_controller=None): - """ Displays a plot, and has the user click on an approximate + """Displays a plot, and has the user click on an approximate zero point. Uses a root finder to adjust to the more exact point. Right click to disable. Parameter: @@ -1017,7 +1101,9 @@ def _toggle_root_finder(self): self._root_finder.toggle_root_finding() def _find_psi_lines(self, tk_controller=None): - self._psi_finder = RootFinder(self.PsiNorm, mode='psi_finder', controller=tk_controller) + self._psi_finder = RootFinder( + self.PsiNorm, mode="psi_finder", controller=tk_controller + ) def GetMagxData(self) -> tuple: """ @@ -1031,7 +1117,11 @@ def GetMagxData(self) -> tuple: ------- A 3-tuple of r-z coordinates and scalar psi value """ - return (self.magx[0], self.magx[1], self.PsiUNorm.get_psi(self.magx[0], self.magx[1])) + return ( + self.magx[0], + self.magx[1], + self.PsiUNorm.get_psi(self.magx[0], self.magx[1]), + ) def GetXptData(self) -> dict: """ @@ -1046,12 +1136,18 @@ def GetXptData(self) -> dict: A dict containing an (r, z, psi) entry for each x-point """ xpt_info = {} - if hasattr(self, 'xpt1'): - xpt_info['xpt1'] = (self.xpt1[0], self.xpt1[1], - self.PsiUNorm.get_psi(self.xpt1[0], self.xpt1[1])) - if hasattr(self, 'xpt2'): - xpt_info['xpt2'] = (self.xpt2[0], self.xpt2[1], - self.PsiUNorm.get_psi(self.xpt2[0], self.xpt2[1])) + if hasattr(self, "xpt1"): + xpt_info["xpt1"] = ( + self.xpt1[0], + self.xpt1[1], + self.PsiUNorm.get_psi(self.xpt1[0], self.xpt1[1]), + ) + if hasattr(self, "xpt2"): + xpt_info["xpt2"] = ( + self.xpt2[0], + self.xpt2[1], + self.PsiUNorm.get_psi(self.xpt2[0], self.xpt2[1]), + ) return xpt_info def PrepLineTracing(self): @@ -1079,13 +1175,13 @@ def GetPatchTagMap(self, config: str) -> dict: A dictionary containing the tag to patch name mappings. """ PatchTagMap = {} - TempLabels = ['A', 'B', 'C', 'D', 'E', 'F'] + TempLabels = ["A", "B", "C", "D", "E", "F"] for label in TempLabels: for i in range(1, 3): PatchTagMap[label + str(i)] = label + str(i) else: PatchTagMap = {} - TempLabels = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I'] + TempLabels = ["A", "B", "C", "D", "E", "F", "G", "H", "I"] for label in TempLabels: for i in range(1, 4): PatchTagMap[label + str(i)] = label + str(i) @@ -1115,35 +1211,46 @@ def ClassifyTopology(self, visual=False) -> None: If user specifies ``settings['grid_settings']['num_xpt']`` with value other than 1 (int) or 2 (int). """ - print('') + print("") print("# Begin classification....") - print('') + print("") - if self.settings['grid_settings']['num_xpt'] == 1: + if self.settings["grid_settings"]["num_xpt"] == 1: self.LineTracer.SNL_find_NSEW(self.xpt1, self.magx, visual) - elif self.settings['grid_settings']['num_xpt'] == 2: + elif self.settings["grid_settings"]["num_xpt"] == 2: # Check if limiter contains magx, xpt1,and xpt2 from matplotlib.patches import Polygon - limiter = Polygon(np.column_stack(self.LimiterData.points()).T, fill=True, closed=True, color='red', label='Limiter') + + limiter = Polygon( + np.column_stack(self.LimiterData.points()).T, + fill=True, + closed=True, + color="red", + label="Limiter", + ) missing_items = [] if (limiter.get_path().contains_point(self.magx)) is False: - missing_items.insert(0, 'magx') + missing_items.insert(0, "magx") if (limiter.get_path().contains_point(self.xpt1)) is False: - missing_items.insert(0, 'xpt1') + missing_items.insert(0, "xpt1") if (limiter.get_path().contains_point(self.xpt2)) is False: - missing_items.insert(0, 'xpt2') + missing_items.insert(0, "xpt2") if len(missing_items) == 0: self.LineTracer.DNL_find_NSEW(self.xpt1, self.xpt2, self.magx, visual) else: - raise ValueError(f"# Topological item(s) {missing_items} not contained in the limiter geometry provided. Check coordinates provided are correct and/or edit the limiter geometry.") + raise ValueError( + f"# Topological item(s) {missing_items} not contained in the limiter geometry provided. Check coordinates provided are correct and/or edit the limiter geometry." + ) else: - raise ValueError(f"# No support available for 'num_xpt' value of {self.settings['grid_settings']['num_xpt']}") + raise ValueError( + f"# No support available for 'num_xpt' value of {self.settings['grid_settings']['num_xpt']}" + ) self.config = self.LineTracer.config - def WriteGridueSNL(self, gridue_settings: dict, fname: str = 'gridue') -> bool: + def WriteGridueSNL(self, gridue_settings: dict, fname: str = "gridue") -> bool: """ Write a gridue file for a single-null configuration. @@ -1162,51 +1269,57 @@ def WriteGridueSNL(self, gridue_settings: dict, fname: str = 'gridue') -> bool: """ def format_header(gridue): - header_items = ['nxm', 'nym', 'ixpt1', 'ixpt2', 'iyseptrx1'] - header = '' + header_items = ["nxm", "nym", "ixpt1", "ixpt2", "iyseptrx1"] + header = "" for item in header_items: - header += '{}'.format(gridue[item]).rjust(4) + header += "{}".format(gridue[item]).rjust(4) - header += '\n' + header += "\n" return header def format_body(data): delim_val = 0 - delim_char = '' - body = '' + delim_char = "" + body = "" for n in range(5): for j in range(len(data[0])): for i in range(len(data)): delim_val += 1 - val = np.format_float_scientific(data[i][j][n], precision=15, unique=False).rjust(23).replace('e', 'D') + val = ( + np.format_float_scientific( + data[i][j][n], precision=15, unique=False + ) + .rjust(23) + .replace("e", "D") + ) if delim_val == 3: delim_val = 0 - delim_char = '\n' + delim_char = "\n" body += val + delim_char - delim_char = '' + delim_char = "" if delim_val % 3 != 0: - body += '\n' + body += "\n" return body - f = open(fname, mode='w') - f.write(format_header(gridue_settings) + '\n') + f = open(fname, mode="w") + f.write(format_header(gridue_settings) + "\n") - body_items = ['rm', 'zm', 'psi', 'br', 'bz', 'bpol', 'bphi', 'b'] + body_items = ["rm", "zm", "psi", "br", "bz", "bpol", "bphi", "b"] for item in body_items: - f.write(format_body(gridue_settings[item]) + '\n') + f.write(format_body(gridue_settings[item]) + "\n") - runidg = 'iogridue' - f.write(runidg + '\n') + runidg = "iogridue" + f.write(runidg + "\n") f.close() return True - def WriteGridueDNL(self, gridue_settings: dict, fname: str = 'gridue') -> bool: + def WriteGridueDNL(self, gridue_settings: dict, fname: str = "gridue") -> bool: """ Write a gridue file for a double-null configuration. @@ -1226,51 +1339,57 @@ def WriteGridueDNL(self, gridue_settings: dict, fname: str = 'gridue') -> bool: def format_header(gridue): header_rows = [ - ['nxm', 'nym'], - ['iyseparatrix1', 'iyseparatrix2'], - ['ix_plate1', 'ix_cut1', '_FILLER_', 'ix_cut2', 'ix_plate2'], - ['iyseparatrix3', 'iyseparatrix4'], - ['ix_plate3', 'ix_cut3', '_FILLER_', 'ix_cut4', 'ix_plate4'] + ["nxm", "nym"], + ["iyseparatrix1", "iyseparatrix2"], + ["ix_plate1", "ix_cut1", "_FILLER_", "ix_cut2", "ix_plate2"], + ["iyseparatrix3", "iyseparatrix4"], + ["ix_plate3", "ix_cut3", "_FILLER_", "ix_cut4", "ix_plate4"], ] - header = '' + header = "" for header_items in header_rows: for item in header_items: - header += '{}'.format(gridue[item]).rjust(4) - header += '\n' + header += "{}".format(gridue[item]).rjust(4) + header += "\n" return header def format_body(data): delim_val = 0 - delim_char = '' - body = '' + delim_char = "" + body = "" for n in range(5): for j in range(len(data[0])): for i in range(len(data)): delim_val += 1 - val = np.format_float_scientific(data[i][j][n], precision=15, unique=False).rjust(23).replace('e', 'D') + val = ( + np.format_float_scientific( + data[i][j][n], precision=15, unique=False + ) + .rjust(23) + .replace("e", "D") + ) if delim_val == 3: delim_val = 0 - delim_char = '\n' + delim_char = "\n" body += val + delim_char - delim_char = '' + delim_char = "" if delim_val % 3 != 0: - body += '\n' + body += "\n" return body - f = open(fname, mode='w') - f.write(format_header(gridue_settings) + '\n') + f = open(fname, mode="w") + f.write(format_header(gridue_settings) + "\n") - body_items = ['rm', 'zm', 'psi', 'br', 'bz', 'bpol', 'bphi', 'b'] + body_items = ["rm", "zm", "psi", "br", "bz", "bpol", "bphi", "b"] for item in body_items: - f.write(format_body(gridue_settings[item]) + '\n') + f.write(format_body(gridue_settings[item]) + "\n") - runidg = 'iogridue' - f.write(runidg + '\n') + runidg = "iogridue" + f.write(runidg + "\n") f.close() return True @@ -1311,10 +1430,13 @@ def ReconstructPatches(self, raw_patch_list: list) -> dict: S = Line([Point(p) for p in zip(SR, SZ)]) W = Line([Point(p) for p in zip(WR, WZ)]) - patch = Patch([N, E, S, W], patch_name=patch_settings['patch_name'], - plate_patch=patch_settings['plate_patch'], - plate_location=patch_settings['plate_location'], - PatchTagMap=patch_settings['PatchTagMap']) + patch = Patch( + [N, E, S, W], + patch_name=patch_settings["patch_name"], + plate_patch=patch_settings["plate_patch"], + plate_location=patch_settings["plate_location"], + PatchTagMap=patch_settings["PatchTagMap"], + ) patches[patch.patch_name] = patch @@ -1350,8 +1472,9 @@ def PrepGridue(self, guard_cell_eps=1e-3) -> None: @classmethod def _CheckOverlapCells(Grid, Verbose=False): from shapely.geometry import Polygon - r = Grid['rm'] - z = Grid['zm'] + + r = Grid["rm"] + z = Grid["zm"] idx = [1, 2, 4, 3] p = [] pinfo = [] @@ -1369,13 +1492,16 @@ def _CheckOverlapCells(Grid, Verbose=False): ListIntersect = [] for p1, pinfo1 in zip(p, pinfo): for p2, pinfo2 in zip(p, pinfo): - if p1.intersects(p2) and np.sum(abs(np.array(pinfo1) - np.array(pinfo2))) > 2: + if ( + p1.intersects(p2) + and np.sum(abs(np.array(pinfo1) - np.array(pinfo2))) > 2 + ): ListIntersect.append((pinfo1, pinfo2)) - print('p1:{} and p2:{} intersect!'.format(pinfo1, pinfo2)) + print("p1:{} and p2:{} intersect!".format(pinfo1, pinfo2)) return ListIntersect -class TopologyUtils(): +class TopologyUtils: """ The base class for all INGRID topologies. @@ -1421,6 +1547,7 @@ class TopologyUtils(): Parent PsiNorm instance. """ + def __init__(self, Ingrid_obj: object, config: str): self.parent = Ingrid_obj self.config = config @@ -1455,30 +1582,45 @@ def patch_diagram(self, fig: object = None, ax: object = None) -> None: """ - colors = {'A': 'red', 'B': 'blue', 'C': 'navajowhite', 'D': 'firebrick', - 'E': 'magenta', 'F': 'olivedrab', 'G': 'darkorange', 'H': 'yellow', 'I': 'navy'} - alpha = {'3': 1.0, '2': 0.75, '1': 0.5} + colors = { + "A": "red", + "B": "blue", + "C": "navajowhite", + "D": "firebrick", + "E": "magenta", + "F": "olivedrab", + "G": "darkorange", + "H": "yellow", + "I": "navy", + } + alpha = {"3": 1.0, "2": 0.75, "1": 0.5} - f = fig if fig else plt.figure('INGRID Patch Map', figsize=(6, 10)) + f = fig if fig else plt.figure("INGRID Patch Map", figsize=(6, 10)) f.subplots_adjust(bottom=0.2) a = ax if ax else f.subplots(1, 1) a.set_xlim([self.PsiUNorm.rmin, self.PsiUNorm.rmax]) a.set_ylim([self.PsiUNorm.zmin, self.PsiUNorm.zmax]) - a.set_aspect('equal', adjustable='box') + a.set_aspect("equal", adjustable="box") - a.set_xlabel('R') - a.set_ylabel('Z') - a.set_title(f'{self.config} Patch Diagram') + a.set_xlabel("R") + a.set_ylabel("Z") + a.set_title(f"{self.config} Patch Diagram") for i, patch in enumerate(self.patches.values()): - patch.plot_border(color='black', ax=a) - patch.fill(colors[patch.get_tag()[0]], ax=a, alpha=alpha[patch.get_tag()[-1]]) + patch.plot_border(color="black", ax=a) + patch.fill( + colors[patch.get_tag()[0]], ax=a, alpha=alpha[patch.get_tag()[-1]] + ) patch.color = colors[patch.get_tag()[0]] handles, labels = a.get_legend_handles_labels() lookup = {label: handle for label, handle in zip(labels, handles)} - a.legend(handles=[handle for handle in lookup.values()], labels=[label for label in lookup.keys()], - bbox_to_anchor=(1.25, 1.0), loc='upper right', - ncol=1) + a.legend( + handles=[handle for handle in lookup.values()], + labels=[label for label in lookup.keys()], + bbox_to_anchor=(1.25, 1.0), + loc="upper right", + ncol=1, + ) f.show() def grid_diagram(self, fig: object = None, ax: object = None) -> None: @@ -1495,7 +1637,7 @@ def grid_diagram(self, fig: object = None, ax: object = None) -> None: """ if fig is None: - fig = plt.figure('INGRID Grid', figsize=(6, 10)) + fig = plt.figure("INGRID Grid", figsize=(6, 10)) if ax is None: ax = fig.gca() @@ -1504,31 +1646,31 @@ def grid_diagram(self, fig: object = None, ax: object = None) -> None: fig.subplots_adjust(bottom=0.2) ax.set_xlim(self.PsiUNorm.rmin, self.PsiUNorm.rmax) ax.set_ylim(self.PsiUNorm.zmin, self.PsiUNorm.zmax) - ax.set_aspect('equal', adjustable='box') + ax.set_aspect("equal", adjustable="box") - ax.set_xlabel('R') - ax.set_ylabel('Z') - ax.set_title(f'{self.config} Subgrid') + ax.set_xlabel("R") + ax.set_ylabel("Z") + ax.set_title(f"{self.config} Subgrid") for patch in self.patches.values(): patch.plot_subgrid(fig=fig, ax=ax) - print(f'# Plotting subgrid {patch.patch_name}') + print(f"# Plotting subgrid {patch.patch_name}") fig.show() def _animate_grid(self): try: - plt.close('INGRID: Debug') + plt.close("INGRID: Debug") except: pass - plt.figure('INGRID: Debug', figsize=(6, 10)) + plt.figure("INGRID: Debug", figsize=(6, 10)) plt.xlim(self.PsiUNorm.rmin, self.PsiUNorm.rmax) plt.ylim(self.PsiUNorm.zmin, self.PsiUNorm.zmax) - plt.gca().set_aspect('equal', adjustable='box') - plt.xlabel('R') - plt.ylabel('Z') - plt.title('visualize gridue') + plt.gca().set_aspect("equal", adjustable="box") + plt.xlabel("R") + plt.ylabel("Z") + plt.title("visualize gridue") k = [1, 2, 4, 3, 1] @@ -1550,7 +1692,7 @@ def get_config(self) -> str: """ return self.config - def get_func(self, func: str) -> 'function': + def get_func(self, func: str) -> "function": """ Create a function from a string input. @@ -1599,16 +1741,17 @@ def get_func(self, func: str) -> 'function': def make_sympy_func(var, expression): import sympy as sp - _f = sp.lambdify(var, expression, 'numpy') + + _f = sp.lambdify(var, expression, "numpy") return _f f_str_raw = func - f_str_raw = f_str_raw.replace(' ', '') - delim = f_str_raw.index(',') + f_str_raw = f_str_raw.replace(" ", "") + delim = f_str_raw.index(",") - var = f_str_raw[0: delim] - expression = f_str_raw[delim + 1:] + var = f_str_raw[0:delim] + expression = f_str_raw[delim + 1 :] func = make_sympy_func(var, expression) # TODO: Check Normalization of the function to 1 @@ -1632,13 +1775,17 @@ def CheckFunction(self, expression: str, Verbose: bool = False) -> bool: """ ExpValid = False try: - com = 'lambda {}:{}'.format(expression.split(',')[0], expression.split(',')[1]) + com = "lambda {}:{}".format( + expression.split(",")[0], expression.split(",")[1] + ) if Verbose: print(com) eval(com) ExpValid = True except: - raise ValueError('Unable to parse expression entry "{}".'.format(expression)) + raise ValueError( + 'Unable to parse expression entry "{}".'.format(expression) + ) return ExpValid def GetFunctions(self, Patch: Patch, Verbose: bool = False) -> tuple: @@ -1664,18 +1811,22 @@ def GetFunctions(self, Patch: Patch, Verbose: bool = False) -> tuple: """ poloidal_tag, radial_tag = Patch.get_tag() - p_f = 'poloidal_f_' + poloidal_tag - r_f = 'radial_f_' + radial_tag + p_f = "poloidal_f_" + poloidal_tag + r_f = "radial_f_" + radial_tag try: - _poloidal_f = self.settings['grid_settings']['grid_generation'][p_f] + _poloidal_f = self.settings["grid_settings"]["grid_generation"][p_f] valid_function = self.CheckFunction(_poloidal_f, Verbose) if valid_function: _poloidal_f = self.get_func(_poloidal_f) else: - raise ValueError('# Invalid function entry. Applying default poloidal function.') + raise ValueError( + "# Invalid function entry. Applying default poloidal function." + ) except: - _poloidal_f = self.settings['grid_settings']['grid_generation']['poloidal_f_default'] + _poloidal_f = self.settings["grid_settings"]["grid_generation"][ + "poloidal_f_default" + ] valid_function = self.CheckFunction(_poloidal_f, Verbose) if valid_function: _poloidal_f = self.get_func(_poloidal_f) @@ -1685,19 +1836,28 @@ def GetFunctions(self, Patch: Patch, Verbose: bool = False) -> tuple: try: # Adding CORE radial_f support for SNL cases via entry 'radial_f_3' - if self.config in ['USN', 'LSN'] \ - and self.settings['grid_settings']['grid_generation'].get('radial_f_3') is not None \ - and poloidal_tag + radial_tag in ['B1', 'C1', 'D1', 'E1']: - _radial_f = self.settings['grid_settings']['grid_generation']['radial_f_3'] + if ( + self.config in ["USN", "LSN"] + and self.settings["grid_settings"]["grid_generation"].get("radial_f_3") + is not None + and poloidal_tag + radial_tag in ["B1", "C1", "D1", "E1"] + ): + _radial_f = self.settings["grid_settings"]["grid_generation"][ + "radial_f_3" + ] else: - _radial_f = self.settings['grid_settings']['grid_generation'][r_f] + _radial_f = self.settings["grid_settings"]["grid_generation"][r_f] valid_function = self.CheckFunction(_radial_f, Verbose) if valid_function: _radial_f = self.get_func(_radial_f) else: - raise ValueError('# Invalid function entry. Applying default radial function.') + raise ValueError( + "# Invalid function entry. Applying default radial function." + ) except: - _radial_f = self.settings['grid_settings']['grid_generation']['radial_f_default'] + _radial_f = self.settings["grid_settings"]["grid_generation"][ + "radial_f_default" + ] valid_function = self.CheckFunction(_radial_f, Verbose) if valid_function: _radial_f = self.get_func(_radial_f) @@ -1724,17 +1884,17 @@ def GetNpoints(self, Patch: Patch) -> tuple: """ poloidal_tag, radial_tag = Patch.get_tag() - np_tag = 'np_' + poloidal_tag - nr_tag = 'nr_' + radial_tag + np_tag = "np_" + poloidal_tag + nr_tag = "nr_" + radial_tag try: - np_cells = self.settings['grid_settings']['grid_generation'][np_tag] + np_cells = self.settings["grid_settings"]["grid_generation"][np_tag] except: - np_cells = self.settings['grid_settings']['grid_generation']['np_default'] + np_cells = self.settings["grid_settings"]["grid_generation"]["np_default"] try: - nr_cells = self.settings['grid_settings']['grid_generation'][nr_tag] + nr_cells = self.settings["grid_settings"]["grid_generation"][nr_tag] except: - nr_cells = self.settings['grid_settings']['grid_generation']['nr_default'] + nr_cells = self.settings["grid_settings"]["grid_generation"]["nr_default"] return (nr_cells, np_cells) @@ -1749,16 +1909,30 @@ def GetDistortionCorrectionSettings(self) -> dict: ------- A dictionary containing ``CorrectDistortion`` settings. """ - if self.settings['grid_settings']['grid_generation'].get('distortion_correction') is not None: - CD = self.settings['grid_settings']['grid_generation']['distortion_correction'] + if ( + self.settings["grid_settings"]["grid_generation"].get( + "distortion_correction" + ) + is not None + ): + CD = self.settings["grid_settings"]["grid_generation"][ + "distortion_correction" + ] else: CD = {} self.distortion_correction = CD return CD - def construct_grid(self, np_cells: int = 1, nr_cells: int = 1, Verbose: bool = False, - ShowVertices: bool = False, RestartScratch: bool = False, ListPatches: str = 'all') -> None: + def construct_grid( + self, + np_cells: int = 1, + nr_cells: int = 1, + Verbose: bool = False, + ShowVertices: bool = False, + RestartScratch: bool = False, + ListPatches: str = "all", + ) -> None: """ Construct a grid by refining a Patch map. @@ -1821,13 +1995,13 @@ def construct_grid(self, np_cells: int = 1, nr_cells: int = 1, Verbose: bool = F # Straighten up East and West segments of our patches, # Plot borders and fill patches. if Verbose: - print('Construct Grid') + print("Construct Grid") try: - visual = self.settings['DEBUG']['visual']['subgrid'] + visual = self.settings["DEBUG"]["visual"]["subgrid"] except: visual = False try: - verbose = self.settings['DEBUG']['verbose']['grid_generation'] + verbose = self.settings["DEBUG"]["verbose"]["grid_generation"] except: verbose = False @@ -1835,7 +2009,7 @@ def construct_grid(self, np_cells: int = 1, nr_cells: int = 1, Verbose: bool = F self.GetDistortionCorrectionSettings() - print('>>> Patches:', [k for k in self.patches.keys()]) + print(">>> Patches:", [k for k in self.patches.keys()]) if RestartScratch: self.CurrentListPatch = {} @@ -1844,24 +2018,37 @@ def construct_grid(self, np_cells: int = 1, nr_cells: int = 1, Verbose: bool = F if self.distortion_correction.get(name) is not None: patch.distortion_correction = self.distortion_correction.get(name) elif self.distortion_correction.get(patch.get_tag()) is not None: - patch.distortion_correction = self.distortion_correction.get(patch.get_tag()) - elif self.distortion_correction.get('all') is not None: - patch.distortion_correction = self.distortion_correction.get('all') + patch.distortion_correction = self.distortion_correction.get( + patch.get_tag() + ) + elif self.distortion_correction.get("all") is not None: + patch.distortion_correction = self.distortion_correction.get("all") else: - patch.distortion_correction = {'Active': False} - if (ListPatches == 'all' and patch not in self.CurrentListPatch) or (ListPatches != 'all' and name in ListPatches): + patch.distortion_correction = {"Active": False} + if (ListPatches == "all" and patch not in self.CurrentListPatch) or ( + ListPatches != "all" and name in ListPatches + ): self.SetPatchBoundaryPoints(patch, verbose) (nr_cells, np_cells) = self.GetNpoints(patch) (_radial_f, _poloidal_f) = self.GetFunctions(patch) - print(f'>>> Making subgrid in patch {name}:') - print(f' np = {np_cells}, nr = {nr_cells}') - print(f' fp = {inspect.getsource(_poloidal_f)}') - print(f' fr = {inspect.getsource(_radial_f)}', end='') + print(f">>> Making subgrid in patch {name}:") + print(f" np = {np_cells}, nr = {nr_cells}") + print(f" fp = {inspect.getsource(_poloidal_f)}") + print(f" fr = {inspect.getsource(_radial_f)}", end="") patch.RemoveDuplicatePoints() - patch.make_subgrid(self, np_cells, nr_cells, _poloidal_f=_poloidal_f, _radial_f=_radial_f, verbose=verbose, visual=visual, ShowVertices=ShowVertices) + patch.make_subgrid( + self, + np_cells, + nr_cells, + _poloidal_f=_poloidal_f, + _radial_f=_radial_f, + verbose=verbose, + visual=visual, + ShowVertices=ShowVertices, + ) self.CurrentListPatch[name] = patch - print(f' {name} subgrid complete.\n\n') + print(f" {name} subgrid complete.\n\n") self.AdjustGrid() def SetPatchBoundaryPoints(self, Patch: Patch, verbose: bool = False) -> None: @@ -1883,13 +2070,13 @@ def SetPatchBoundaryPoints(self, Patch: Patch, verbose: bool = False) -> None: Patch.TerminatesLoop = False if self.ConnexionMap.get(Patch.get_tag()) is not None: if verbose: - print('Find connexion map for patch {}'.format(Patch.patch_name)) + print("Find connexion map for patch {}".format(Patch.patch_name)) for v in self.ConnexionMap.get(Patch.get_tag()).items(): Boundary, AdjacentPatch = v Patch.BoundaryPoints[Boundary] = self.GetBoundaryPoints(AdjacentPatch) if verbose: - print('Find Boundaries points for {}'.format(Patch.patch_name)) - if self.ConnexionMap.get(Patch.get_tag()).get('E') is not None: + print("Find Boundaries points for {}".format(Patch.patch_name)) + if self.ConnexionMap.get(Patch.get_tag()).get("E") is not None: Patch.TerminatesLoop = True def GetBoundaryPoints(self, AdjacentPatchInfo: tuple) -> list: @@ -1912,13 +2099,13 @@ def GetBoundaryPoints(self, AdjacentPatchInfo: tuple) -> list: Boundary = AdjacentPatchInfo[1] for patch in self.patches.values(): if patch.get_tag() == PatchTag: - if Boundary == 'S': + if Boundary == "S": return patch.S_vertices - elif Boundary == 'N': + elif Boundary == "N": return patch.N_vertices - elif Boundary == 'E': + elif Boundary == "E": return patch.E_vertices - elif Boundary == 'W': + elif Boundary == "W": return patch.W_vertices return None @@ -1936,7 +2123,7 @@ def CheckPatches(self, verbose: bool = False) -> None: """ for name, patch in self.patches.items(): if patch.plate_patch: - print(' # Checking patch: ', name) + print(" # Checking patch: ", name) patch.CheckPatch(self) def SetupPatchMatrix(self) -> list: @@ -1955,38 +2142,252 @@ def SetupPatchMatrix(self) -> list: """ p = self.patches - if self.config in ['LSN', 'USN']: - self.patch_matrix = [ - [[None], [None], [None], [None], [None], [None], [None], [None]], - [[None], p['A2'], p['B2'], p['C2'], p['D2'], p['E2'], p['F2'], [None]], - [[None], p['A1'], p['B1'], p['C1'], p['D1'], p['E1'], p['F1'], [None]], - [[None], [None], [None], [None], [None], [None], [None], [None]] - ] - - elif self.config in ['SF45', 'SF75', 'SF105', 'SF135']: + if self.config in ["LSN", "USN"]: + if self.settings["grid_settings"]["up_down_symmetry"]: + #TODO: Implement up/down symmetry for other geometries + self.patch_matrix = [ + [[None], [None], [None], [None], [None], [None], [None], [None]], + [[None], p["A2"], p["B2"], [None], [None], p["E2"], p["F2"], [None]], + [[None], p["A1"], p["B1"], [None], [None], p["E1"], p["F1"], [None]], + [[None], [None], [None], [None], [None], [None], [None], [None]], + ] + else: + self.patch_matrix = [ + [[None], [None], [None], [None], [None], [None], [None], [None]], + [[None], p['A2'], p['B2'], p['C2'], p['D2'], p['E2'], p['F2'], [None]], + [[None], p['A1'], p['B1'], p['C1'], p['D1'], p['E1'], p['F1'], [None]], + [[None], [None], [None], [None], [None], [None], [None], [None]] + ] + + elif self.config in ["SF45", "SF75", "SF105", "SF135"]: self.patch_matrix = [ - [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]], - [[None], p['A3'], p['B3'], p['C3'], p['D3'], p['E3'], p['F3'], p['G3'], [None], [None], p['H3'], p['I3'], [None]], - [[None], p['A2'], p['B2'], p['C2'], p['D2'], p['E2'], p['F2'], p['G2'], [None], [None], p['H2'], p['I2'], [None]], - [[None], p['A1'], p['B1'], p['C1'], p['D1'], p['E1'], p['F1'], p['G1'], [None], [None], p['H1'], p['I1'], [None]], - [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]] + [ + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + ], + [ + [None], + p["A3"], + p["B3"], + p["C3"], + p["D3"], + p["E3"], + p["F3"], + p["G3"], + [None], + [None], + p["H3"], + p["I3"], + [None], + ], + [ + [None], + p["A2"], + p["B2"], + p["C2"], + p["D2"], + p["E2"], + p["F2"], + p["G2"], + [None], + [None], + p["H2"], + p["I2"], + [None], + ], + [ + [None], + p["A1"], + p["B1"], + p["C1"], + p["D1"], + p["E1"], + p["F1"], + p["G1"], + [None], + [None], + p["H1"], + p["I1"], + [None], + ], + [ + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + ], ] - elif self.config in ['SF15', 'SF165']: + elif self.config in ["SF15", "SF165"]: self.patch_matrix = [ - [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]], - [[None], p['A3'], p['B3'], p['C3'], p['D3'], p['E3'], p['F3'], [None], [None], p['G3'], p['H3'], p['I3'], [None]], - [[None], p['A2'], p['B2'], p['C2'], p['D2'], p['E2'], p['F2'], [None], [None], p['G2'], p['H2'], p['I2'], [None]], - [[None], p['A1'], p['B1'], p['C1'], p['D1'], p['E1'], p['F1'], [None], [None], p['G1'], p['H1'], p['I1'], [None]], - [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]] + [ + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + ], + [ + [None], + p["A3"], + p["B3"], + p["C3"], + p["D3"], + p["E3"], + p["F3"], + [None], + [None], + p["G3"], + p["H3"], + p["I3"], + [None], + ], + [ + [None], + p["A2"], + p["B2"], + p["C2"], + p["D2"], + p["E2"], + p["F2"], + [None], + [None], + p["G2"], + p["H2"], + p["I2"], + [None], + ], + [ + [None], + p["A1"], + p["B1"], + p["C1"], + p["D1"], + p["E1"], + p["F1"], + [None], + [None], + p["G1"], + p["H1"], + p["I1"], + [None], + ], + [ + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + ], ] - elif self.config in ['UDN']: + elif self.config in ["UDN"]: self.patch_matrix = [ - [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]], - [[None], p['A3'], p['B3'], p['C3'], p['D3'], [None], [None], p['E3'], p['F3'], p['G3'], p['H3'], [None]], - [[None], p['A2'], p['B2'], p['C2'], p['D2'], [None], [None], p['E2'], p['F2'], p['G2'], p['H2'], [None]], - [[None], p['A1'], p['B1'], p['C1'], p['D1'], [None], [None], p['E1'], p['F1'], p['G1'], p['H1'], [None]], - [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]] + [ + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + ], + [ + [None], + p["A3"], + p["B3"], + p["C3"], + p["D3"], + [None], + [None], + p["E3"], + p["F3"], + p["G3"], + p["H3"], + [None], + ], + [ + [None], + p["A2"], + p["B2"], + p["C2"], + p["D2"], + [None], + [None], + p["E2"], + p["F2"], + p["G2"], + p["H2"], + [None], + ], + [ + [None], + p["A1"], + p["B1"], + p["C1"], + p["D1"], + [None], + [None], + p["E1"], + p["F1"], + p["G1"], + p["H1"], + [None], + ], + [ + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + [None], + ], ] return self.patch_matrix @@ -2014,40 +2415,76 @@ def _add_guardc(cell_map, ixlb, ixrb, nxpt=1, eps=1e-3): def set_guard(cell_map, ix, iy, eps, boundary): # Note: 'USN' and 'right' is really just 'LSN' and 'left' settings. # TODO: Edit the code to reflect this at some point so the next reader is not overwhelmed. - if boundary == 'left': + if boundary == "left": ixn = ix + 1 iyn = iy - cell_map[ix][iy][1] = cell_map[ixn][iyn][1] + eps * (cell_map[ixn][iyn][1] - cell_map[ixn][iyn][2]) + cell_map[ix][iy][1] = cell_map[ixn][iyn][1] + eps * ( + cell_map[ixn][iyn][1] - cell_map[ixn][iyn][2] + ) cell_map[ix][iy][2] = cell_map[ixn][iyn][1] - cell_map[ix][iy][3] = cell_map[ixn][iyn][3] + eps * (cell_map[ixn][iyn][3] - cell_map[ixn][iyn][4]) + cell_map[ix][iy][3] = cell_map[ixn][iyn][3] + eps * ( + cell_map[ixn][iyn][3] - cell_map[ixn][iyn][4] + ) cell_map[ix][iy][4] = cell_map[ixn][iyn][3] - cell_map[ix][iy][0] = 0.25 * (cell_map[ix][iy][1] + cell_map[ix][iy][2] + cell_map[ix][iy][3] + cell_map[ix][iy][4]) - - elif boundary == 'right': + cell_map[ix][iy][0] = 0.25 * ( + cell_map[ix][iy][1] + + cell_map[ix][iy][2] + + cell_map[ix][iy][3] + + cell_map[ix][iy][4] + ) + + elif boundary == "right": ixn = ix - 1 iyn = iy - cell_map[ix][iy][2] = cell_map[ixn][iyn][2] + eps * (cell_map[ixn][iyn][2] - cell_map[ixn][iyn][1]) + cell_map[ix][iy][2] = cell_map[ixn][iyn][2] + eps * ( + cell_map[ixn][iyn][2] - cell_map[ixn][iyn][1] + ) cell_map[ix][iy][1] = cell_map[ixn][iyn][2] - cell_map[ix][iy][4] = cell_map[ixn][iyn][4] + eps * (cell_map[ixn][iyn][4] - cell_map[ixn][iyn][3]) + cell_map[ix][iy][4] = cell_map[ixn][iyn][4] + eps * ( + cell_map[ixn][iyn][4] - cell_map[ixn][iyn][3] + ) cell_map[ix][iy][3] = cell_map[ixn][iyn][4] - cell_map[ix][iy][0] = 0.25 * (cell_map[ix][iy][1] + cell_map[ix][iy][2] + cell_map[ix][iy][3] + cell_map[ix][iy][4]) - - elif boundary == 'bottom': + cell_map[ix][iy][0] = 0.25 * ( + cell_map[ix][iy][1] + + cell_map[ix][iy][2] + + cell_map[ix][iy][3] + + cell_map[ix][iy][4] + ) + + elif boundary == "bottom": ixn = ix iyn = iy + 1 - cell_map[ix][iy][1] = cell_map[ixn][iyn][1] + eps * (cell_map[ixn][iyn][1] - cell_map[ixn][iyn][3]) + cell_map[ix][iy][1] = cell_map[ixn][iyn][1] + eps * ( + cell_map[ixn][iyn][1] - cell_map[ixn][iyn][3] + ) cell_map[ix][iy][3] = cell_map[ixn][iyn][1] - cell_map[ix][iy][2] = cell_map[ixn][iyn][2] + eps * (cell_map[ixn][iyn][2] - cell_map[ixn][iyn][4]) + cell_map[ix][iy][2] = cell_map[ixn][iyn][2] + eps * ( + cell_map[ixn][iyn][2] - cell_map[ixn][iyn][4] + ) cell_map[ix][iy][4] = cell_map[ixn][iyn][2] - cell_map[ix][iy][0] = 0.25 * (cell_map[ix][iy][1] + cell_map[ix][iy][2] + cell_map[ix][iy][3] + cell_map[ix][iy][4]) - elif boundary == 'top': + cell_map[ix][iy][0] = 0.25 * ( + cell_map[ix][iy][1] + + cell_map[ix][iy][2] + + cell_map[ix][iy][3] + + cell_map[ix][iy][4] + ) + elif boundary == "top": ixn = ix iyn = iy - 1 - cell_map[ix][iy][3] = cell_map[ixn][iyn][3] + eps * (cell_map[ixn][iyn][3] - cell_map[ixn][iyn][1]) + cell_map[ix][iy][3] = cell_map[ixn][iyn][3] + eps * ( + cell_map[ixn][iyn][3] - cell_map[ixn][iyn][1] + ) cell_map[ix][iy][1] = cell_map[ixn][iyn][3] - cell_map[ix][iy][4] = cell_map[ixn][iyn][4] + eps * (cell_map[ixn][iyn][4] - cell_map[ixn][iyn][2]) + cell_map[ix][iy][4] = cell_map[ixn][iyn][4] + eps * ( + cell_map[ixn][iyn][4] - cell_map[ixn][iyn][2] + ) cell_map[ix][iy][2] = cell_map[ixn][iyn][4] - cell_map[ix][iy][0] = 0.25 * (cell_map[ix][iy][1] + cell_map[ix][iy][2] + cell_map[ix][iy][3] + cell_map[ix][iy][4]) + cell_map[ix][iy][0] = 0.25 * ( + cell_map[ix][iy][1] + + cell_map[ix][iy][2] + + cell_map[ix][iy][3] + + cell_map[ix][iy][4] + ) return cell_map @@ -2056,15 +2493,15 @@ def set_guard(cell_map, ix, iy, eps, boundary): for iy in range(1, nr + 1): ix = ixlb - cell_map = set_guard(cell_map, ix, iy, eps, boundary='left') + cell_map = set_guard(cell_map, ix, iy, eps, boundary="left") ix = ixrb + 1 - cell_map = set_guard(cell_map, ix, iy, eps, boundary='right') + cell_map = set_guard(cell_map, ix, iy, eps, boundary="right") for ix in range(np + 2): iy = 0 - cell_map = set_guard(cell_map, ix, iy, eps, boundary='bottom') + cell_map = set_guard(cell_map, ix, iy, eps, boundary="bottom") iy = nr + 1 - cell_map = set_guard(cell_map, ix, iy, eps, boundary='top') + cell_map = set_guard(cell_map, ix, iy, eps, boundary="top") return cell_map @@ -2074,7 +2511,7 @@ def set_guard(cell_map, ix, iy, eps, boundary): patch.npol = len(patch.cell_grid[0]) + 1 patch.nrad = len(patch.cell_grid) + 1 - if self.parent.settings['grid_settings']['num_xpt'] == 1: + if (self.parent.settings['grid_settings']['num_xpt'] == 1) and (self.parent.settings['grid_settings']["up_down_symmetry"] is False): np_total = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][1:-1]])) + 2 nr_total = int(np.sum([patch[1].nrad - 1 for patch in patch_matrix[1:3]])) + 2 @@ -2131,37 +2568,58 @@ def set_guard(cell_map, ix, iy, eps, boundary): if debug: self._animate_grid() - elif self.parent.settings['grid_settings']['num_xpt'] == 2: + elif (self.parent.settings['grid_settings']['num_xpt'] == 2) or ((self.parent.settings['grid_settings']['num_xpt'] == 1) and (self.parent.settings['grid_settings']["up_down_symmetry"] is True)): - if self.config in ['SF45', 'SF75', 'SF105', 'SF135']: + if self.config in ["SF45", "SF75", "SF105", "SF135"]: pindex1 = 8 pindex2 = 10 pindex3 = 12 - elif self.config in ['SF15', 'SF165']: + elif self.config in ["SF15", "SF165"]: pindex1 = 7 pindex2 = 9 pindex3 = 12 - elif self.config in ['UDN']: + elif self.config in ["UDN"]: pindex1 = 5 pindex2 = 7 pindex3 = 11 + elif self.config in ["LSN", "USN"] and (self.parent.settings['grid_settings']["up_down_symmetry"] is True): + pindex1 = 3 + pindex2 = 5 + pindex3 = 7 # Total number of poloidal indices in all subgrids. - np_total1 = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][1:pindex1]])) + 2 + np_total1 = ( + int(np.sum([patch.npol - 1 for patch in patch_matrix[1][1:pindex1]])) + + 2 + ) # Total number of radial indices in all subgrids. - nr_total1 = int(np.sum([patch[1].nrad - 1 for patch in patch_matrix[1:4]])) + 2 + # nr_total1 = int(np.sum([patch[1].nrad - 1 for patch in patch_matrix[1:4]])) + 2 + nr_total1 = ( + int(np.sum([patch[1].nrad - 1 for patch in patch_matrix[1:3]])) + 2 + ) # Total number of poloidal indices in all subgrids. - np_total2 = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][pindex2:pindex3]])) + 2 + np_total2 = ( + int( + np.sum( + [patch.npol - 1 for patch in patch_matrix[1][pindex2:pindex3]] + ) + ) + + 2 + ) # Total number of radial indices in all subgrids. - nr_total2 = int(np.sum([patch[pindex2].nrad - 1 for patch in patch_matrix[1:4]])) + 2 + # nr_total2 = int(np.sum([patch[pindex2].nrad - 1 for patch in patch_matrix[1:4]])) + 2 + nr_total2 = ( + int(np.sum([patch[pindex2].nrad - 1 for patch in patch_matrix[1:3]])) + + 2 + ) - rm1 = np.zeros((np_total1, nr_total1, 5), order='F') - zm1 = np.zeros((np_total1, nr_total1, 5), order='F') - rm2 = np.zeros((np_total2, nr_total2, 5), order='F') - zm2 = np.zeros((np_total2, nr_total2, 5), order='F') + rm1 = np.zeros((np_total1, nr_total1, 5), order="F") + zm1 = np.zeros((np_total1, nr_total1, 5), order="F") + rm2 = np.zeros((np_total2, nr_total2, 5), order="F") + zm2 = np.zeros((np_total2, nr_total2, 5), order="F") ixcell = 0 jycell = 0 @@ -2185,15 +2643,30 @@ def set_guard(cell_map, ix, iy, eps, boundary): # jyl - number of radial cells in the patch for jyl in range(len(local_patch.cell_grid)): - ixcell = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][1:ixp + 1]])) \ - - len(local_patch.cell_grid[0]) + ixl + 1 + ixcell = ( + int( + np.sum( + [ + patch.npol - 1 + for patch in patch_matrix[1][1 : ixp + 1] + ] + ) + ) + - len(local_patch.cell_grid[0]) + + ixl + + 1 + ) jycell = nr_sum - (local_patch.nrad - 1) + jyl + 1 ind = 0 - for coor in ['CENTER', 'SW', 'SE', 'NW', 'NE']: - rm1[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].x - zm1[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].y + for coor in ["CENTER", "SW", "SE", "NW", "NE"]: + rm1[ixcell][jycell][ind] = ( + local_patch.cell_grid[jyl][ixl].vertices[coor].x + ) + zm1[ixcell][jycell][ind] = ( + local_patch.cell_grid[jyl][ixl].vertices[coor].y + ) ind += 1 ixcell = 0 @@ -2217,15 +2690,32 @@ def set_guard(cell_map, ix, iy, eps, boundary): # jyl - number of radial cells in the patch for jyl in range(len(local_patch.cell_grid)): - ixcell = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][pindex2:ixp + 1]])) \ - - len(local_patch.cell_grid[0]) + ixl + 1 + ixcell = ( + int( + np.sum( + [ + patch.npol - 1 + for patch in patch_matrix[1][ + pindex2 : ixp + 1 + ] + ] + ) + ) + - len(local_patch.cell_grid[0]) + + ixl + + 1 + ) jycell = nr_sum - (local_patch.nrad - 1) + jyl + 1 ind = 0 - for coor in ['CENTER', 'SW', 'SE', 'NW', 'NE']: - rm2[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].x - zm2[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].y + for coor in ["CENTER", "SW", "SE", "NW", "NE"]: + rm2[ixcell][jycell][ind] = ( + local_patch.cell_grid[jyl][ixl].vertices[coor].x + ) + zm2[ixcell][jycell][ind] = ( + local_patch.cell_grid[jyl][ixl].vertices[coor].y + ) ind += 1 # Flip indices into gridue format. @@ -2253,7 +2743,7 @@ def set_guard(cell_map, ix, iy, eps, boundary): self.zm = np.concatenate((zm1, zm2)) try: - debug = self.yaml['DEBUG']['visual']['gridue'] + debug = self.yaml["DEBUG"]["visual"]["gridue"] except: debug = False From 084e96c7ca584b3a81ce58bbad83498df94349ab Mon Sep 17 00:00:00 2001 From: Dominic Power Date: Tue, 6 May 2025 14:25:07 -0400 Subject: [PATCH 02/10] Bug fix in export grid and indid formatting changes --- INGRID/utils.py | 1406 ++++++++++++++++------------------------------- 1 file changed, 467 insertions(+), 939 deletions(-) diff --git a/INGRID/utils.py b/INGRID/utils.py index 518df7d..99754c1 100644 --- a/INGRID/utils.py +++ b/INGRID/utils.py @@ -10,7 +10,6 @@ from __future__ import print_function, division, absolute_import import numpy as np import matplotlib - try: matplotlib.use("TkAgg") except: @@ -27,7 +26,7 @@ from INGRID.geometry import Point, Line, Patch, orientation_between -class IngridUtils: +class IngridUtils(): """ The base class for :class:`ingrid.Ingrid` that handles backend management of key Ingrid capabilities. This class can be directly utilized by advanced users @@ -120,11 +119,11 @@ def __init__(self, settings: dict = {}, **kwargs): self.InputFile = None self.config = None self.settings_lookup = [ - "dir_settings", - "grid_settings", - "integrator_settings", - "target_plates", - "DEBUG", + 'dir_settings', + 'grid_settings', + 'integrator_settings', + 'target_plates', + 'DEBUG' ] self.SetDefaultSettings() self.PopulateSettings(settings) @@ -161,125 +160,148 @@ def SetDefaultSettings(self) -> None: """ self.default_grid_settings = { - "num_xpt": 1, - "nlevs": 30, - "view_mode": "filled", - "psi_1": 0.0, - "psi_core": 0.0, - "psi_pf_1": 0.0, - "psi_pf_2": 0.0, - "psi_1": 0.0, - "psi_2": 0.0, - "rmagx": 0.0, - "zmagx": 0.0, - "rxpt": 0.0, - "zxpt": 0.0, - "rxpt2": 0.0, - "zxpt2": 0.0, - "guard_cell_eps": 1.0e-3, - "grid_generation": { - "np_default": 2, - "nr_default": 2, - "poloidal_f_default": "x, x", - "radial_f_default": "x, x", - "distortion_correction": { - "all": { - "theta_min": 80.0, - "theta_max": 120.0, - "resolution": 1000, - "active": False, + 'num_xpt': 1, + 'nlevs': 30, + 'view_mode': 'filled', + 'psi_1': 0.0, + 'psi_core': 0.0, + 'psi_pf_1': 0.0, + 'psi_pf_2': 0.0, + 'psi_1': 0.0, + 'psi_2': 0.0, + 'rmagx': 0.0, + 'zmagx': 0.0, + 'rxpt': 0.0, + 'zxpt': 0.0, + 'rxpt2': 0.0, + 'zxpt2': 0.0, + 'guard_cell_eps': 1.0e-3, + 'grid_generation': { + 'np_default': 2, + 'nr_default': 2, + 'poloidal_f_default': 'x, x', + 'radial_f_default': 'x, x', + 'distortion_correction': { + 'all': { + 'theta_min': 80.0, + 'theta_max': 120.0, + 'resolution': 1000, + 'active': False }, }, }, - "patch_generation": { - "core_split_point_ratio": 0.5, - "pf_split_point_ratio": 0.5, - "strike_pt_loc": "limiter", - "rmagx_shift": 0.0, - "zmagx_shift": 0.0, - "magx_tilt_1": 0.0, - "magx_tilt_2": 0.0, - "use_xpt1_W": False, - "use_xpt1_E": False, - "use_xpt2_W": False, - "use_xpt2_E": False, - "xpt1_W_tilt": -0.785398, # All values of pi / 4 radians. - "xpt1_E_tilt": 0.785398, - "xpt2_W_tilt": -0.785398, - "xpt2_E_tilt": 0.785398, + 'patch_generation': { + 'core_split_point_ratio': 0.5, + 'pf_split_point_ratio': 0.5, + 'strike_pt_loc': 'limiter', + 'rmagx_shift': 0.0, + 'zmagx_shift': 0.0, + 'magx_tilt_1': 0.0, + 'magx_tilt_2': 0.0, + 'use_xpt1_W': False, + 'use_xpt1_E': False, + 'use_xpt2_W': False, + 'use_xpt2_E': False, + 'xpt1_W_tilt': -0.785398, # All values of pi / 4 radians. + 'xpt1_E_tilt': 0.785398, + 'xpt2_W_tilt': -0.785398, + 'xpt2_E_tilt': 0.785398, }, "up_down_symmetry": False } self.default_integrator_settings = { - "dt": 0.01, - "eps": 5e-5, - "first_step": 1e-5, - "step_ratio": 0.02, - "tol": 5e-3, - "max_step": 0.064, + 'dt': 0.01, + 'eps': 5e-5, + 'first_step': 1e-5, + 'step_ratio': 0.02, + 'tol': 5e-3, + 'max_step': 0.064 } self.default_target_plate_settings = { - "plate_E1": {"file": "", "rshift": 0.0, "zshift": 0.0}, - "plate_E2": {"file": "", "rshift": 0.0, "zshift": 0.0}, - "plate_W1": {"file": "", "rshift": 0.0, "zshift": 0.0}, - "plate_W2": {"file": "", "rshift": 0.0, "zshift": 0.0}, + 'plate_E1': { + 'file': '', + 'rshift': 0.0, + 'zshift': 0.0 + }, + + 'plate_E2': { + 'file': '', + 'rshift': 0.0, + 'zshift': 0.0 + }, + + 'plate_W1': { + 'file': '', + 'rshift': 0.0, + 'zshift': 0.0 + }, + + 'plate_W2': { + 'file': '', + 'rshift': 0.0, + 'zshift': 0.0 + }, } self.default_limiter_settings = { - "file": "", - "use_efit_bounds": False, - "efit_buffer_r": 1e-2, - "efit_buffer_z": 1e-2, - "rshift": 0.0, - "zshift": 0.0, + 'file': '', + 'use_efit_bounds': False, + 'efit_buffer_r': 1e-2, + 'efit_buffer_z': 1e-2, + 'rshift': 0.0, + 'zshift': 0.0 } self.default_patch_data_settings = { - "file": "", - "use_file": False, - "preferences": {"new_file": False, "new_fname": ""}, + 'file': '', + 'use_file': False, + 'preferences': { + 'new_file': False, + 'new_fname': '' + } } self.default_DEBUG_settings = { - "visual": { - "find_NSEW": False, - "patch_map": False, - "subgrid": False, - "gridue": False, - "SF_analysis": False, - }, - "verbose": { - "patch_generation": False, - "grid_generation": False, - "SF_analysis": False, + 'visual': { + 'find_NSEW': False, + 'patch_map': False, + 'subgrid': False, + 'gridue': False, + 'SF_analysis': False }, + + 'verbose': { + 'patch_generation': False, + 'grid_generation': False, + 'SF_analysis': False + } } self.default_dir_settings = { - "eqdsk": ".", - "limiter": ".", - "patch_data": ".", - "target_plates": ".", + 'eqdsk': '.', + 'limiter': '.', + 'patch_data': '.', + 'target_plates': '.' } self.default_values_lookup = { - "eqdsk": "", - "dir_settings": self.default_dir_settings, - "grid_settings": self.default_grid_settings, - "integrator_settings": self.default_integrator_settings, - "target_plates": self.default_target_plate_settings, - "limiter": self.default_limiter_settings, - "patch_data": self.default_patch_data_settings, - "DEBUG": self.default_DEBUG_settings, + 'eqdsk': '', + 'dir_settings': self.default_dir_settings, + 'grid_settings': self.default_grid_settings, + 'integrator_settings': self.default_integrator_settings, + 'target_plates': self.default_target_plate_settings, + 'limiter': self. default_limiter_settings, + 'patch_data': self.default_patch_data_settings, + 'DEBUG': self.default_DEBUG_settings } self.PlateData = { - "plate_W1": {}, - "plate_E1": {}, - "plate_W2": {}, - "plate_E2": {}, + 'plate_W1': {}, + 'plate_E1': {}, + 'plate_W2': {}, + 'plate_E2': {} } def ProcessKeywords(self, **kwargs) -> None: @@ -287,27 +309,27 @@ def ProcessKeywords(self, **kwargs) -> None: Process kwargs and set all file paths accordingly. """ for k, v in kwargs.items(): - if k == "InputFile" or k == "yaml": - print("# Processing Input File:", v) + if k == 'InputFile' or k == 'yaml': + print('# Processing Input File:', v) self.InputFile = v self.PopulateSettings(self.ReadYamlFile(v)) continue - if k == "W1TargetFile" or k == "w1": - self.settings["target_plates"]["plate_W1"]["file"] = v + if k == 'W1TargetFile' or k == 'w1': + self.settings['target_plates']['plate_W1']['file'] = v continue - if k == "E1TargetFile" or k == "e1": - self.settings["target_plates"]["plate_E1"]["file"] = v + if k == 'E1TargetFile' or k == 'e1': + self.settings['target_plates']['plate_E1']['file'] = v continue - if k == "E2TargetFile" or k == "e2": - self.settings["target_plates"]["plate_E2"]["file"] = v + if k == 'E2TargetFile' or k == 'e2': + self.settings['target_plates']['plate_E2']['file'] = v continue - if k == "W2TargetFile" or k == "w2": - self.settings["target_plates"]["plate_W2"]["file"] = v + if k == 'W2TargetFile' or k == 'w2': + self.settings['target_plates']['plate_W2']['file'] = v continue - if k in ["LimiterFile", "limiter", "wall"]: - self.settings["limiter"]["file"] = v - if k == "EqFile" or k == "eq": - self.settings["eqdsk"] = v + if k in ['LimiterFile', 'limiter', 'wall']: + self.settings['limiter']['file'] = v + if k == 'EqFile' or k == 'eq': + self.settings['eqdsk'] = v continue print('Keyword "' + k + '" unknown and ignored...') @@ -337,9 +359,7 @@ def _check_settings_input(input_settings, comparison): raise_assertion = True if raise_assertion is True: - raise ValueError( - f"Invalid entries {items} in provided settings file f{self.InputFile}. Remove invalid entries listed." - ) + raise ValueError(f'Invalid entries {items} in provided settings file f{self.InputFile}. Remove invalid entries listed.') def _get_default_values(item, sub_item=None, attribute=None): """ @@ -350,11 +370,7 @@ def _get_default_values(item, sub_item=None, attribute=None): try: default_values = self.default_values_lookup[item] except KeyError: - print( - f"Key " - "{item}" - " not recognized... Add default values to source code for support." - ) + print(f'Key ''{item}'' not recognized... Add default values to source code for support.') return None if item and sub_item and attribute: @@ -364,9 +380,7 @@ def _get_default_values(item, sub_item=None, attribute=None): elif item: return self.default_values_lookup[item] - _check_settings_input( - input_settings=settings, comparison=self.default_values_lookup - ) + _check_settings_input(input_settings=settings, comparison=self.default_values_lookup) # First level entries within YAML file dump. for item in self.default_values_lookup.keys(): @@ -378,20 +392,13 @@ def _get_default_values(item, sub_item=None, attribute=None): settings[item] = _get_default_values(item) if verbose: print('Could not find "{}" in YAML file.'.format(item)) - print( - 'Populated "{}" with default value of "{}".\n'.format( - item, settings[item] - ) - ) + print('Populated "{}" with default value of "{}".\n'.format(item, settings[item])) continue - if item in ["eqdsk"]: + if item in ['eqdsk']: continue - _check_settings_input( - input_settings=settings[item], - comparison=self.default_values_lookup[item], - ) + _check_settings_input(input_settings=settings[item], comparison=self.default_values_lookup[item]) # Second level entries within YAML file dump. @@ -401,45 +408,25 @@ def _get_default_values(item, sub_item=None, attribute=None): except KeyError: settings[item][sub_item] = _get_default_values(item, sub_item) if verbose: - print( - 'Could not find "{}/{}" in YAML file.'.format( - item, sub_item - ) - ) - print( - 'Populated "{}/{}" with default value of "{}".\n'.format( - item, sub_item, settings[item][sub_item] - ) - ) + print('Could not find "{}/{}" in YAML file.'.format(item, sub_item)) + print('Populated "{}/{}" with default value of "{}".\n'.format(item, sub_item, settings[item][sub_item])) continue - if item in ["grid_settings", "target_plates"] and sub_item in [ - "patch_generation", - "grid_generation", - "plate_E1", - "plate_W1", - "plate_E2", - "plate_W2", - ]: - for plate_attribute in self.default_values_lookup[item][ - sub_item - ].keys(): + if item in ['grid_settings', 'target_plates'] \ + and sub_item in ['patch_generation', 'grid_generation', 'plate_E1', 'plate_W1', 'plate_E2', 'plate_W2']: + for plate_attribute in self.default_values_lookup[item][sub_item].keys(): try: if settings[item][sub_item][plate_attribute] is None: - settings[item][sub_item][plate_attribute] = ( - _get_default_values(item, sub_item, plate_attribute) - ) + settings[item][sub_item][plate_attribute] = _get_default_values(item, sub_item, plate_attribute) except: - settings[item][sub_item][plate_attribute] = ( - _get_default_values(item, sub_item, plate_attribute) - ) + settings[item][sub_item][plate_attribute] = _get_default_values(item, sub_item, plate_attribute) # Update references to YAML entries. self.settings = settings - self.grid_settings = settings["grid_settings"] - self.integrator_settings = settings["integrator_settings"] - self.target_plates = settings["target_plates"] - self.DEBUG = settings["DEBUG"] + self.grid_settings = settings['grid_settings'] + self.integrator_settings = settings['integrator_settings'] + self.target_plates = settings['target_plates'] + self.DEBUG = settings['DEBUG'] self.ProcessPaths() def ProcessPaths(self) -> None: @@ -447,55 +434,26 @@ def ProcessPaths(self) -> None: Update settings by pre-pending path entries to all file entries. """ for k in self.default_dir_settings.keys(): - path_obj = Path(self.settings["dir_settings"][k]) + path_obj = Path(self.settings['dir_settings'][k]) if path_obj.is_dir() is False: - v = self.settings["dir_settings"][k] - raise ValueError( - f"# Error processing directory provided for entry '{k}'. Entry '{v}' is not a valid directory." - ) - - if k == "eqdsk": - self.settings["eqdsk"] = str( - (path_obj / self.settings["eqdsk"]).absolute() - ) + v = self.settings['dir_settings'][k] + raise ValueError(f"# Error processing directory provided for entry '{k}'. Entry '{v}' is not a valid directory.") + + if k == 'eqdsk': + self.settings['eqdsk'] = str((path_obj / self.settings['eqdsk']).absolute()) continue - if k == "limiter": - self.settings["limiter"]["file"] = str( - (path_obj / self.settings["limiter"]["file"]).absolute() - ) + if k == 'limiter': + self.settings['limiter']['file'] = str((path_obj / self.settings['limiter']['file']).absolute()) continue - if k == "target_plates": - self.settings["target_plates"]["plate_W1"]["file"] = str( - ( - path_obj / self.settings["target_plates"]["plate_W1"]["file"] - ).absolute() - ) - self.settings["target_plates"]["plate_E1"]["file"] = str( - ( - path_obj / self.settings["target_plates"]["plate_E1"]["file"] - ).absolute() - ) - self.settings["target_plates"]["plate_W2"]["file"] = str( - ( - path_obj / self.settings["target_plates"]["plate_W2"]["file"] - ).absolute() - ) - self.settings["target_plates"]["plate_E2"]["file"] = str( - ( - path_obj / self.settings["target_plates"]["plate_E2"]["file"] - ).absolute() - ) + if k == 'target_plates': + self.settings['target_plates']['plate_W1']['file'] = str((path_obj / self.settings['target_plates']['plate_W1']['file']).absolute()) + self.settings['target_plates']['plate_E1']['file'] = str((path_obj / self.settings['target_plates']['plate_E1']['file']).absolute()) + self.settings['target_plates']['plate_W2']['file'] = str((path_obj / self.settings['target_plates']['plate_W2']['file']).absolute()) + self.settings['target_plates']['plate_E2']['file'] = str((path_obj / self.settings['target_plates']['plate_E2']['file']).absolute()) continue - if k == "patch_data": - self.settings["patch_data"]["file"] = str( - (path_obj / self.settings["patch_data"]["file"]).absolute() - ) - self.settings["patch_data"]["preferences"]["new_fname"] = str( - ( - path_obj - / self.settings["patch_data"]["preferences"]["new_fname"] - ).absolute() - ) + if k == 'patch_data': + self.settings['patch_data']['file'] = str((path_obj / self.settings['patch_data']['file']).absolute()) + self.settings['patch_data']['preferences']['new_fname'] = str((path_obj / self.settings['patch_data']['preferences']['new_fname']).absolute()) continue def LoadGEQDSK(self, geqdsk_path: str) -> None: @@ -505,66 +463,59 @@ def LoadGEQDSK(self, geqdsk_path: str) -> None: Saves the boundary information and generates an EfitData instance. """ - with open(geqdsk_path, "r") as f: + with open(geqdsk_path, 'r') as f: geqdsk_data = geqdsk.read(f) # # Extract quantities needed to initialize EfitData class # - nx = geqdsk_data["nx"] - ny = geqdsk_data["ny"] - rdim = geqdsk_data["rdim"] - zdim = geqdsk_data["zdim"] - zmid = geqdsk_data["zmid"] - rleft = geqdsk_data["rleft"] - rcenter = geqdsk_data["rcentr"] - bcenter = geqdsk_data["bcentr"] - rmagx = geqdsk_data["rmagx"] - zmagx = geqdsk_data["zmagx"] - rlimiter = geqdsk_data["rlim"] - zlimiter = geqdsk_data["zlim"] - psi = geqdsk_data["psi"] + nx = geqdsk_data['nx'] + ny = geqdsk_data['ny'] + rdim = geqdsk_data['rdim'] + zdim = geqdsk_data['zdim'] + zmid = geqdsk_data['zmid'] + rleft = geqdsk_data['rleft'] + rcenter = geqdsk_data['rcentr'] + bcenter = geqdsk_data['bcentr'] + rmagx = geqdsk_data['rmagx'] + zmagx = geqdsk_data['zmagx'] + rlimiter = geqdsk_data['rlim'] + zlimiter = geqdsk_data['zlim'] + psi = geqdsk_data['psi'] # # Derived values # rmin = rleft rmax = rmin + rdim - zmin = zmid - 0.5 * zdim + zmin = (zmid - 0.5 * zdim) zmax = zmin + zdim # # Reproduce efit grid # self.PsiUNorm = EfitData( - rmin=rmin, - rmax=rmax, - nr=nx, - zmin=zmin, - zmax=zmax, - nz=ny, - rcenter=rcenter, - bcenter=bcenter, - rlimiter=rlimiter, - zlimiter=zlimiter, - rmagx=rmagx, - zmagx=zmagx, - name="Efit Data", - parent=self, - ) - - self.PsiUNorm.init_bivariate_spline( - self.PsiUNorm.r[:, 0], self.PsiUNorm.z[0, :], psi - ) - - if ( - self.settings["grid_settings"]["rmagx"] is None - or self.settings["grid_settings"]["zmagx"] is None - ): - ( - self.settings["grid_settings"]["rmagx"], - self.settings["grid_settings"]["zmagx"], - ) = (self.PsiUNorm.rmagx, self.PsiUNorm.zmagx) + rmin=rmin, + rmax=rmax, + nr=nx, + zmin=zmin, + zmax=zmax, + nz=ny, + rcenter=rcenter, + bcenter=bcenter, + rlimiter=rlimiter, + zlimiter=zlimiter, + rmagx=rmagx, + zmagx=zmagx, + name='Efit Data', + parent=self) + + self.PsiUNorm.init_bivariate_spline(self.PsiUNorm.r[:, 0], + self.PsiUNorm.z[0, :], + psi) + + if self.settings['grid_settings']['rmagx'] is None or self.settings['grid_settings']['zmagx'] is None: + self.settings['grid_settings']['rmagx'], self.settings['grid_settings']['zmagx'] = (self.PsiUNorm.rmagx, self.PsiUNorm.zmagx) # # Save the stored data in the object @@ -573,11 +524,10 @@ def LoadGEQDSK(self, geqdsk_path: str) -> None: # # Update the eqdsk file referenced in settings to that of the loaded data # - self.settings["eqdsk"] = geqdsk_path + self.settings['eqdsk'] = geqdsk_path + - def ParseTxtCoordinates( - self, fpath: str, rshift: float = 0.0, zshift: float = 0.0 - ) -> tuple: + def ParseTxtCoordinates(self, fpath: str, rshift: float = 0.0, zshift: float = 0.0) -> tuple: """ Extract the (R,Z) coordinates from a .txt file. @@ -629,52 +579,37 @@ def ParseTxtCoordinates( R, Z = [], [] - if Path(fpath).is_file() and Path(fpath).suffix in [".txt"]: + if Path(fpath).is_file() and Path(fpath).suffix in ['.txt']: try: with open(fpath) as f: for line in f: - if line.startswith("#"): + if line.startswith('#'): # move along to the next iteration # this simulates a comment continue # we deal with separator ',' or ' ' - point = ( - line.replace("D", "e") - .replace("(", "") - .replace(")", "") - .strip() - ) - if point.count(",") > 0: - x = float(point.split(",")[0]) - y = float(point.split(",")[1]) + point = line.replace('D', 'e').replace('(', '').replace(')', '').strip() + if point.count(',') > 0: + x = float(point.split(',')[0]) + y = float(point.split(',')[1]) else: - x = float(point.split(" ")[0]) - y = float(point.split(" ")[1]) + x = float(point.split(' ')[0]) + y = float(point.split(' ')[1]) R.append(x + rshift) Z.append(y + zshift) except: - raise IOError( - f"# Error occured when reading data from file {fpath}:\t 'open(fpath)' error" - ) + raise IOError(f"# Error occured when reading data from file {fpath}:\t 'open(fpath)' error") else: - raise ValueError( - f"# Error occur when reading data from file {fpath}:\t file does not exist or is not of extension '*.txt'" - ) + raise ValueError(f"# Error occur when reading data from file {fpath}:\t file does not exist or is not of extension '*.txt'") return R, Z - def SetLimiter( - self, - fpath: str = "", - coordinates: list = [], - rshift: float = 0.0, - zshift: float = 0.0, - ) -> None: + def SetLimiter(self, fpath: str = '', coordinates: list = [], rshift: float = 0.0, zshift: float = 0.0) -> None: """ Instantiate the class Line object that represents the tokamak limiter. @@ -701,30 +636,28 @@ def SetLimiter( """ if rshift is None: - rshift = self.settings["limiter"]["rshift"] + rshift = self.settings['limiter']['rshift'] if zshift is None: - zshift = self.settings["limiter"]["zshift"] + zshift = self.settings['limiter']['zshift'] - use_efit_bounds = self.settings["limiter"]["use_efit_bounds"] + use_efit_bounds = self.settings['limiter']['use_efit_bounds'] # # TODO: Refactor the codebase to leverage numpy arrays.... # coordinates = np.array(coordinates).T - + # # If we have no coordinates or fpath provided, we need to # leverage efit boundaries. Overrides any user settings. # - if coordinates.shape[0] == 0 and fpath == "": + if coordinates.shape[0] == 0 and fpath == '': use_efit_bounds = True - if fpath not in ["", "."]: + if fpath not in ['', '.']: try: - print("# Processing file for limiter data : {}".format(fpath)) - self.geqdsk_data["rlim"], self.geqdsk_data["rlim"] = ( - self.ParseTxtCoordinates(fpath) - ) + print('# Processing file for limiter data : {}'.format(fpath)) + self.geqdsk_data['rlim'], self.geqdsk_data['rlim'] = self.ParseTxtCoordinates(fpath) except: raise ValueError(f"# Error in method 'SetLimiter' with fpath={fpath}") @@ -732,16 +665,16 @@ def SetLimiter( # List of coordinates provided for initialization # elif len(coordinates) > 0: - self.geqdsk_data["rlim"] = coordinates[:, 0] + rshift - self.geqdsk_data["zlim"] = coordinates[:, 1] + zshift + self.geqdsk_data['rlim'] = coordinates[:, 0] + rshift + self.geqdsk_data['zlim'] = coordinates[:, 1] + zshift # # Empty list of coordinates falls back on using eqdsk limiter settings # else: - self.LoadGEQDSK(geqdsk_path=self.settings["eqdsk"]) - self.geqdsk_data["rlim"] += rshift - self.geqdsk_data["zlim"] += zshift + self.LoadGEQDSK(geqdsk_path=self.settings['eqdsk']) + self.geqdsk_data['rlim'] += rshift + self.geqdsk_data['zlim'] += zshift if use_efit_bounds: @@ -750,8 +683,8 @@ def SetLimiter( # Buffer values shrink the bounding box to allow us to avoid stepping # outside the domain during integration/line-tracing # - rbuff = self.settings["limiter"]["efit_buffer_r"] - zbuff = self.settings["limiter"]["efit_buffer_z"] + rbuff = self.settings['limiter']['efit_buffer_r'] + zbuff = self.settings['limiter']['efit_buffer_z'] # # Define the bounding box vertices as: @@ -768,17 +701,13 @@ def SetLimiter( # # Define the simple limiter path and translate by shift values # - limiter_path = np.vstack([LL, LR, UR, UL, LL]) - self.geqdsk_data["rlim"] = limiter_path[:, 0] + rshift - self.geqdsk_data["zlim"] = limiter_path[:, 1] + zshift + limiter_path = np.vstack([LL, LR, UR, UL, LL]) + self.geqdsk_data['rlim'] = limiter_path[:, 0] + rshift + self.geqdsk_data['zlim'] = limiter_path[:, 1] + zshift - self.LimiterData = Line( - [Point(p) for p in zip(self.geqdsk_data["rlim"], self.geqdsk_data["zlim"])] - ) + self.LimiterData = Line([Point(p) for p in zip(self.geqdsk_data['rlim'], self.geqdsk_data['zlim'])]) - def SetTargetPlate( - self, settings: dict, rshift: float = 0.0, zshift: float = 0.0 - ) -> None: + def SetTargetPlate(self, settings: dict, rshift: float = 0.0, zshift: float = 0.0) -> None: """ Initialize a target plate Line object. @@ -844,14 +773,14 @@ def SetTargetPlate( v = _v break - if k.lower() in ["w1", "westplate1", "plate_w1"]: - k = "plate_W1" - elif k.lower() in ["e1", "eastplate1", "plate_e1"]: - k = "plate_E1" - elif k.lower() in ["w2", "westplate2", "plate_w2"]: - k = "plate_W2" - elif k.lower() in ["e2", "eastplate2", "plate_e2"]: - k = "plate_E2" + if k.lower() in ['w1', 'westplate1', 'plate_w1']: + k = 'plate_W1' + elif k.lower() in ['e1', 'eastplate1', 'plate_e1']: + k = 'plate_E1' + elif k.lower() in ['w2', 'westplate2', 'plate_w2']: + k = 'plate_W2' + elif k.lower() in ['e2', 'eastplate2', 'plate_e2']: + k = 'plate_E2' else: raise ValueError(f"# Invalid key '{k}' provided for 'SetTargetPlate'") @@ -863,9 +792,7 @@ def SetTargetPlate( data = np.array([c for c in zip(R, Z)]) a, index = np.unique(data, return_index=True, axis=0) index.sort() - self.PlateData[k] = Line( - [Point(x + rshift, y + zshift) for x, y in data[index]] - ) + self.PlateData[k] = Line([Point(x + rshift, y + zshift) for x, y in data[index]]) def OrderTargetPlate(self, plate_key: str) -> None: """ @@ -897,9 +824,7 @@ def OrderTargetPlate(self, plate_key: str) -> None: k = plate_key if not isinstance(self.PlateData.get(k), Line): - raise ValueError( - f"# Plate '{k}' is not loaded as a Line object. Make sure 'SetGeometry({{'{k}' : ... }})' has been called." - ) + raise ValueError(f"# Plate '{k}' is not loaded as a Line object. Make sure 'SetGeometry({{'{k}' : ... }})' has been called.") plate = self.PlateData[k] @@ -930,12 +855,8 @@ def OrderTargetPlate(self, plate_key: str) -> None: # Check the angle between the plate endpoints with # tail fixed on the magnetic axis - v_reference = np.array( - [ - self.settings["grid_settings"]["rmagx"], - self.settings["grid_settings"]["zmagx"], - ] - ) + v_reference = np.array([self.settings['grid_settings']['rmagx'], + self.settings['grid_settings']['zmagx']]) v_start = np.array([start.x, start.y]) v_end = np.array([end.x, end.y]) @@ -1016,12 +937,8 @@ def OrderLimiter(self) -> None: # Check the angle between the plate endpoints with # tail fixed on the magnetic axis - v_reference = np.array( - [ - self.settings["grid_settings"]["rmagx"], - self.settings["grid_settings"]["zmagx"], - ] - ) + v_reference = np.array([self.settings['grid_settings']['rmagx'], + self.settings['grid_settings']['zmagx']]) v_start = np.array([start.x, start.y]) v_end = np.array([end.x, end.y]) @@ -1045,8 +962,8 @@ def FindMagAxis(self, r: float, z: float) -> None: Z coordinate of magnetic-axis guess. """ sol = root(self.PsiUNorm.Gradient, [r, z]) - self.settings["grid_settings"]["rmagx"] = sol.x[0] - self.settings["grid_settings"]["zmagx"] = sol.x[1] + self.settings['grid_settings']['rmagx'] = sol.x[0] + self.settings['grid_settings']['zmagx'] = sol.x[1] self.magx = (sol.x[0], sol.x[1]) def FindXPoint(self, r: float, z: float) -> None: @@ -1062,8 +979,8 @@ def FindXPoint(self, r: float, z: float) -> None: Z coordinate of primary x-point guess. """ sol = root(self.PsiUNorm.Gradient, [r, z]) - self.settings["grid_settings"]["rxpt"] = sol.x[0] - self.settings["grid_settings"]["zxpt"] = sol.x[1] + self.settings['grid_settings']['rxpt'] = sol.x[0] + self.settings['grid_settings']['zxpt'] = sol.x[1] self.xpt1 = (sol.x[0], sol.x[1]) def FindXPoint2(self, r: float, z: float) -> None: @@ -1079,12 +996,12 @@ def FindXPoint2(self, r: float, z: float) -> None: Z coordinate of secondary x-point guess. """ sol = root(self.PsiUNorm.Gradient, [r, z]) - self.settings["grid_settings"]["rxpt2"] = sol.x[0] - self.settings["grid_settings"]["zxpt2"] = sol.x[1] + self.settings['grid_settings']['rxpt2'] = sol.x[0] + self.settings['grid_settings']['zxpt2'] = sol.x[1] self.xpt2 = (sol.x[0], sol.x[1]) def _find_roots(self, tk_controller=None): - """Displays a plot, and has the user click on an approximate + """ Displays a plot, and has the user click on an approximate zero point. Uses a root finder to adjust to the more exact point. Right click to disable. Parameter: @@ -1101,9 +1018,7 @@ def _toggle_root_finder(self): self._root_finder.toggle_root_finding() def _find_psi_lines(self, tk_controller=None): - self._psi_finder = RootFinder( - self.PsiNorm, mode="psi_finder", controller=tk_controller - ) + self._psi_finder = RootFinder(self.PsiNorm, mode='psi_finder', controller=tk_controller) def GetMagxData(self) -> tuple: """ @@ -1117,11 +1032,7 @@ def GetMagxData(self) -> tuple: ------- A 3-tuple of r-z coordinates and scalar psi value """ - return ( - self.magx[0], - self.magx[1], - self.PsiUNorm.get_psi(self.magx[0], self.magx[1]), - ) + return (self.magx[0], self.magx[1], self.PsiUNorm.get_psi(self.magx[0], self.magx[1])) def GetXptData(self) -> dict: """ @@ -1136,18 +1047,12 @@ def GetXptData(self) -> dict: A dict containing an (r, z, psi) entry for each x-point """ xpt_info = {} - if hasattr(self, "xpt1"): - xpt_info["xpt1"] = ( - self.xpt1[0], - self.xpt1[1], - self.PsiUNorm.get_psi(self.xpt1[0], self.xpt1[1]), - ) - if hasattr(self, "xpt2"): - xpt_info["xpt2"] = ( - self.xpt2[0], - self.xpt2[1], - self.PsiUNorm.get_psi(self.xpt2[0], self.xpt2[1]), - ) + if hasattr(self, 'xpt1'): + xpt_info['xpt1'] = (self.xpt1[0], self.xpt1[1], + self.PsiUNorm.get_psi(self.xpt1[0], self.xpt1[1])) + if hasattr(self, 'xpt2'): + xpt_info['xpt2'] = (self.xpt2[0], self.xpt2[1], + self.PsiUNorm.get_psi(self.xpt2[0], self.xpt2[1])) return xpt_info def PrepLineTracing(self): @@ -1175,13 +1080,13 @@ def GetPatchTagMap(self, config: str) -> dict: A dictionary containing the tag to patch name mappings. """ PatchTagMap = {} - TempLabels = ["A", "B", "C", "D", "E", "F"] + TempLabels = ['A', 'B', 'C', 'D', 'E', 'F'] for label in TempLabels: for i in range(1, 3): PatchTagMap[label + str(i)] = label + str(i) else: PatchTagMap = {} - TempLabels = ["A", "B", "C", "D", "E", "F", "G", "H", "I"] + TempLabels = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I'] for label in TempLabels: for i in range(1, 4): PatchTagMap[label + str(i)] = label + str(i) @@ -1211,46 +1116,35 @@ def ClassifyTopology(self, visual=False) -> None: If user specifies ``settings['grid_settings']['num_xpt']`` with value other than 1 (int) or 2 (int). """ - print("") + print('') print("# Begin classification....") - print("") + print('') - if self.settings["grid_settings"]["num_xpt"] == 1: + if self.settings['grid_settings']['num_xpt'] == 1: self.LineTracer.SNL_find_NSEW(self.xpt1, self.magx, visual) - elif self.settings["grid_settings"]["num_xpt"] == 2: + elif self.settings['grid_settings']['num_xpt'] == 2: # Check if limiter contains magx, xpt1,and xpt2 from matplotlib.patches import Polygon - - limiter = Polygon( - np.column_stack(self.LimiterData.points()).T, - fill=True, - closed=True, - color="red", - label="Limiter", - ) + limiter = Polygon(np.column_stack(self.LimiterData.points()).T, fill=True, closed=True, color='red', label='Limiter') missing_items = [] if (limiter.get_path().contains_point(self.magx)) is False: - missing_items.insert(0, "magx") + missing_items.insert(0, 'magx') if (limiter.get_path().contains_point(self.xpt1)) is False: - missing_items.insert(0, "xpt1") + missing_items.insert(0, 'xpt1') if (limiter.get_path().contains_point(self.xpt2)) is False: - missing_items.insert(0, "xpt2") + missing_items.insert(0, 'xpt2') if len(missing_items) == 0: self.LineTracer.DNL_find_NSEW(self.xpt1, self.xpt2, self.magx, visual) else: - raise ValueError( - f"# Topological item(s) {missing_items} not contained in the limiter geometry provided. Check coordinates provided are correct and/or edit the limiter geometry." - ) + raise ValueError(f"# Topological item(s) {missing_items} not contained in the limiter geometry provided. Check coordinates provided are correct and/or edit the limiter geometry.") else: - raise ValueError( - f"# No support available for 'num_xpt' value of {self.settings['grid_settings']['num_xpt']}" - ) + raise ValueError(f"# No support available for 'num_xpt' value of {self.settings['grid_settings']['num_xpt']}") self.config = self.LineTracer.config - def WriteGridueSNL(self, gridue_settings: dict, fname: str = "gridue") -> bool: + def WriteGridueSNL(self, gridue_settings: dict, fname: str = 'gridue') -> bool: """ Write a gridue file for a single-null configuration. @@ -1269,57 +1163,51 @@ def WriteGridueSNL(self, gridue_settings: dict, fname: str = "gridue") -> bool: """ def format_header(gridue): - header_items = ["nxm", "nym", "ixpt1", "ixpt2", "iyseptrx1"] - header = "" + header_items = ['nxm', 'nym', 'ixpt1', 'ixpt2', 'iyseptrx1'] + header = '' for item in header_items: - header += "{}".format(gridue[item]).rjust(4) + header += '{}'.format(gridue[item]).rjust(4) - header += "\n" + header += '\n' return header def format_body(data): delim_val = 0 - delim_char = "" - body = "" + delim_char = '' + body = '' for n in range(5): for j in range(len(data[0])): for i in range(len(data)): delim_val += 1 - val = ( - np.format_float_scientific( - data[i][j][n], precision=15, unique=False - ) - .rjust(23) - .replace("e", "D") - ) + val = np.format_float_scientific(data[i][j][n], precision=15, unique=False).rjust(23).replace('e', 'D') if delim_val == 3: delim_val = 0 - delim_char = "\n" + delim_char = '\n' body += val + delim_char - delim_char = "" + delim_char = '' if delim_val % 3 != 0: - body += "\n" + body += '\n' return body - f = open(fname, mode="w") - f.write(format_header(gridue_settings) + "\n") + f = open(fname, mode='w') + f.write(format_header(gridue_settings) + '\n') - body_items = ["rm", "zm", "psi", "br", "bz", "bpol", "bphi", "b"] + body_items = ['rm', 'zm', 'psi', 'br', 'bz', 'bpol', 'bphi', 'b'] for item in body_items: - f.write(format_body(gridue_settings[item]) + "\n") + f.write(format_body(gridue_settings[item]) + '\n') - runidg = "iogridue" - f.write(runidg + "\n") + runidg = 'iogridue' + f.write(runidg + '\n') f.close() return True - def WriteGridueDNL(self, gridue_settings: dict, fname: str = "gridue") -> bool: + def WriteGridueDNL(self, gridue_settings: dict, fname: str = 'gridue') -> bool: """ Write a gridue file for a double-null configuration. @@ -1339,57 +1227,51 @@ def WriteGridueDNL(self, gridue_settings: dict, fname: str = "gridue") -> bool: def format_header(gridue): header_rows = [ - ["nxm", "nym"], - ["iyseparatrix1", "iyseparatrix2"], - ["ix_plate1", "ix_cut1", "_FILLER_", "ix_cut2", "ix_plate2"], - ["iyseparatrix3", "iyseparatrix4"], - ["ix_plate3", "ix_cut3", "_FILLER_", "ix_cut4", "ix_plate4"], + ['nxm', 'nym'], + ['iyseparatrix1', 'iyseparatrix2'], + ['ix_plate1', 'ix_cut1', '_FILLER_', 'ix_cut2', 'ix_plate2'], + ['iyseparatrix3', 'iyseparatrix4'], + ['ix_plate3', 'ix_cut3', '_FILLER_', 'ix_cut4', 'ix_plate4'] ] - header = "" + header = '' for header_items in header_rows: for item in header_items: - header += "{}".format(gridue[item]).rjust(4) - header += "\n" + header += '{}'.format(gridue[item]).rjust(4) + header += '\n' return header def format_body(data): delim_val = 0 - delim_char = "" - body = "" + delim_char = '' + body = '' for n in range(5): for j in range(len(data[0])): for i in range(len(data)): delim_val += 1 - val = ( - np.format_float_scientific( - data[i][j][n], precision=15, unique=False - ) - .rjust(23) - .replace("e", "D") - ) + val = np.format_float_scientific(data[i][j][n], precision=15, unique=False).rjust(23).replace('e', 'D') if delim_val == 3: delim_val = 0 - delim_char = "\n" + delim_char = '\n' body += val + delim_char - delim_char = "" + delim_char = '' if delim_val % 3 != 0: - body += "\n" + body += '\n' return body - f = open(fname, mode="w") - f.write(format_header(gridue_settings) + "\n") + f = open(fname, mode='w') + f.write(format_header(gridue_settings) + '\n') - body_items = ["rm", "zm", "psi", "br", "bz", "bpol", "bphi", "b"] + body_items = ['rm', 'zm', 'psi', 'br', 'bz', 'bpol', 'bphi', 'b'] for item in body_items: - f.write(format_body(gridue_settings[item]) + "\n") + f.write(format_body(gridue_settings[item]) + '\n') - runidg = "iogridue" - f.write(runidg + "\n") + runidg = 'iogridue' + f.write(runidg + '\n') f.close() return True @@ -1430,13 +1312,10 @@ def ReconstructPatches(self, raw_patch_list: list) -> dict: S = Line([Point(p) for p in zip(SR, SZ)]) W = Line([Point(p) for p in zip(WR, WZ)]) - patch = Patch( - [N, E, S, W], - patch_name=patch_settings["patch_name"], - plate_patch=patch_settings["plate_patch"], - plate_location=patch_settings["plate_location"], - PatchTagMap=patch_settings["PatchTagMap"], - ) + patch = Patch([N, E, S, W], patch_name=patch_settings['patch_name'], + plate_patch=patch_settings['plate_patch'], + plate_location=patch_settings['plate_location'], + PatchTagMap=patch_settings['PatchTagMap']) patches[patch.patch_name] = patch @@ -1472,9 +1351,8 @@ def PrepGridue(self, guard_cell_eps=1e-3) -> None: @classmethod def _CheckOverlapCells(Grid, Verbose=False): from shapely.geometry import Polygon - - r = Grid["rm"] - z = Grid["zm"] + r = Grid['rm'] + z = Grid['zm'] idx = [1, 2, 4, 3] p = [] pinfo = [] @@ -1492,16 +1370,13 @@ def _CheckOverlapCells(Grid, Verbose=False): ListIntersect = [] for p1, pinfo1 in zip(p, pinfo): for p2, pinfo2 in zip(p, pinfo): - if ( - p1.intersects(p2) - and np.sum(abs(np.array(pinfo1) - np.array(pinfo2))) > 2 - ): + if p1.intersects(p2) and np.sum(abs(np.array(pinfo1) - np.array(pinfo2))) > 2: ListIntersect.append((pinfo1, pinfo2)) - print("p1:{} and p2:{} intersect!".format(pinfo1, pinfo2)) + print('p1:{} and p2:{} intersect!'.format(pinfo1, pinfo2)) return ListIntersect -class TopologyUtils: +class TopologyUtils(): """ The base class for all INGRID topologies. @@ -1547,7 +1422,6 @@ class TopologyUtils: Parent PsiNorm instance. """ - def __init__(self, Ingrid_obj: object, config: str): self.parent = Ingrid_obj self.config = config @@ -1582,45 +1456,30 @@ def patch_diagram(self, fig: object = None, ax: object = None) -> None: """ - colors = { - "A": "red", - "B": "blue", - "C": "navajowhite", - "D": "firebrick", - "E": "magenta", - "F": "olivedrab", - "G": "darkorange", - "H": "yellow", - "I": "navy", - } - alpha = {"3": 1.0, "2": 0.75, "1": 0.5} + colors = {'A': 'red', 'B': 'blue', 'C': 'navajowhite', 'D': 'firebrick', + 'E': 'magenta', 'F': 'olivedrab', 'G': 'darkorange', 'H': 'yellow', 'I': 'navy'} + alpha = {'3': 1.0, '2': 0.75, '1': 0.5} - f = fig if fig else plt.figure("INGRID Patch Map", figsize=(6, 10)) + f = fig if fig else plt.figure('INGRID Patch Map', figsize=(6, 10)) f.subplots_adjust(bottom=0.2) a = ax if ax else f.subplots(1, 1) a.set_xlim([self.PsiUNorm.rmin, self.PsiUNorm.rmax]) a.set_ylim([self.PsiUNorm.zmin, self.PsiUNorm.zmax]) - a.set_aspect("equal", adjustable="box") + a.set_aspect('equal', adjustable='box') - a.set_xlabel("R") - a.set_ylabel("Z") - a.set_title(f"{self.config} Patch Diagram") + a.set_xlabel('R') + a.set_ylabel('Z') + a.set_title(f'{self.config} Patch Diagram') for i, patch in enumerate(self.patches.values()): - patch.plot_border(color="black", ax=a) - patch.fill( - colors[patch.get_tag()[0]], ax=a, alpha=alpha[patch.get_tag()[-1]] - ) + patch.plot_border(color='black', ax=a) + patch.fill(colors[patch.get_tag()[0]], ax=a, alpha=alpha[patch.get_tag()[-1]]) patch.color = colors[patch.get_tag()[0]] handles, labels = a.get_legend_handles_labels() lookup = {label: handle for label, handle in zip(labels, handles)} - a.legend( - handles=[handle for handle in lookup.values()], - labels=[label for label in lookup.keys()], - bbox_to_anchor=(1.25, 1.0), - loc="upper right", - ncol=1, - ) + a.legend(handles=[handle for handle in lookup.values()], labels=[label for label in lookup.keys()], + bbox_to_anchor=(1.25, 1.0), loc='upper right', + ncol=1) f.show() def grid_diagram(self, fig: object = None, ax: object = None) -> None: @@ -1637,7 +1496,7 @@ def grid_diagram(self, fig: object = None, ax: object = None) -> None: """ if fig is None: - fig = plt.figure("INGRID Grid", figsize=(6, 10)) + fig = plt.figure('INGRID Grid', figsize=(6, 10)) if ax is None: ax = fig.gca() @@ -1646,31 +1505,31 @@ def grid_diagram(self, fig: object = None, ax: object = None) -> None: fig.subplots_adjust(bottom=0.2) ax.set_xlim(self.PsiUNorm.rmin, self.PsiUNorm.rmax) ax.set_ylim(self.PsiUNorm.zmin, self.PsiUNorm.zmax) - ax.set_aspect("equal", adjustable="box") + ax.set_aspect('equal', adjustable='box') - ax.set_xlabel("R") - ax.set_ylabel("Z") - ax.set_title(f"{self.config} Subgrid") + ax.set_xlabel('R') + ax.set_ylabel('Z') + ax.set_title(f'{self.config} Subgrid') for patch in self.patches.values(): patch.plot_subgrid(fig=fig, ax=ax) - print(f"# Plotting subgrid {patch.patch_name}") + print(f'# Plotting subgrid {patch.patch_name}') fig.show() def _animate_grid(self): try: - plt.close("INGRID: Debug") + plt.close('INGRID: Debug') except: pass - plt.figure("INGRID: Debug", figsize=(6, 10)) + plt.figure('INGRID: Debug', figsize=(6, 10)) plt.xlim(self.PsiUNorm.rmin, self.PsiUNorm.rmax) plt.ylim(self.PsiUNorm.zmin, self.PsiUNorm.zmax) - plt.gca().set_aspect("equal", adjustable="box") - plt.xlabel("R") - plt.ylabel("Z") - plt.title("visualize gridue") + plt.gca().set_aspect('equal', adjustable='box') + plt.xlabel('R') + plt.ylabel('Z') + plt.title('visualize gridue') k = [1, 2, 4, 3, 1] @@ -1692,7 +1551,7 @@ def get_config(self) -> str: """ return self.config - def get_func(self, func: str) -> "function": + def get_func(self, func: str) -> 'function': """ Create a function from a string input. @@ -1741,17 +1600,16 @@ def get_func(self, func: str) -> "function": def make_sympy_func(var, expression): import sympy as sp - - _f = sp.lambdify(var, expression, "numpy") + _f = sp.lambdify(var, expression, 'numpy') return _f f_str_raw = func - f_str_raw = f_str_raw.replace(" ", "") - delim = f_str_raw.index(",") + f_str_raw = f_str_raw.replace(' ', '') + delim = f_str_raw.index(',') - var = f_str_raw[0:delim] - expression = f_str_raw[delim + 1 :] + var = f_str_raw[0: delim] + expression = f_str_raw[delim + 1:] func = make_sympy_func(var, expression) # TODO: Check Normalization of the function to 1 @@ -1775,17 +1633,13 @@ def CheckFunction(self, expression: str, Verbose: bool = False) -> bool: """ ExpValid = False try: - com = "lambda {}:{}".format( - expression.split(",")[0], expression.split(",")[1] - ) + com = 'lambda {}:{}'.format(expression.split(',')[0], expression.split(',')[1]) if Verbose: print(com) eval(com) ExpValid = True except: - raise ValueError( - 'Unable to parse expression entry "{}".'.format(expression) - ) + raise ValueError('Unable to parse expression entry "{}".'.format(expression)) return ExpValid def GetFunctions(self, Patch: Patch, Verbose: bool = False) -> tuple: @@ -1811,22 +1665,18 @@ def GetFunctions(self, Patch: Patch, Verbose: bool = False) -> tuple: """ poloidal_tag, radial_tag = Patch.get_tag() - p_f = "poloidal_f_" + poloidal_tag - r_f = "radial_f_" + radial_tag + p_f = 'poloidal_f_' + poloidal_tag + r_f = 'radial_f_' + radial_tag try: - _poloidal_f = self.settings["grid_settings"]["grid_generation"][p_f] + _poloidal_f = self.settings['grid_settings']['grid_generation'][p_f] valid_function = self.CheckFunction(_poloidal_f, Verbose) if valid_function: _poloidal_f = self.get_func(_poloidal_f) else: - raise ValueError( - "# Invalid function entry. Applying default poloidal function." - ) + raise ValueError('# Invalid function entry. Applying default poloidal function.') except: - _poloidal_f = self.settings["grid_settings"]["grid_generation"][ - "poloidal_f_default" - ] + _poloidal_f = self.settings['grid_settings']['grid_generation']['poloidal_f_default'] valid_function = self.CheckFunction(_poloidal_f, Verbose) if valid_function: _poloidal_f = self.get_func(_poloidal_f) @@ -1836,28 +1686,19 @@ def GetFunctions(self, Patch: Patch, Verbose: bool = False) -> tuple: try: # Adding CORE radial_f support for SNL cases via entry 'radial_f_3' - if ( - self.config in ["USN", "LSN"] - and self.settings["grid_settings"]["grid_generation"].get("radial_f_3") - is not None - and poloidal_tag + radial_tag in ["B1", "C1", "D1", "E1"] - ): - _radial_f = self.settings["grid_settings"]["grid_generation"][ - "radial_f_3" - ] + if self.config in ['USN', 'LSN'] \ + and self.settings['grid_settings']['grid_generation'].get('radial_f_3') is not None \ + and poloidal_tag + radial_tag in ['B1', 'C1', 'D1', 'E1']: + _radial_f = self.settings['grid_settings']['grid_generation']['radial_f_3'] else: - _radial_f = self.settings["grid_settings"]["grid_generation"][r_f] + _radial_f = self.settings['grid_settings']['grid_generation'][r_f] valid_function = self.CheckFunction(_radial_f, Verbose) if valid_function: _radial_f = self.get_func(_radial_f) else: - raise ValueError( - "# Invalid function entry. Applying default radial function." - ) + raise ValueError('# Invalid function entry. Applying default radial function.') except: - _radial_f = self.settings["grid_settings"]["grid_generation"][ - "radial_f_default" - ] + _radial_f = self.settings['grid_settings']['grid_generation']['radial_f_default'] valid_function = self.CheckFunction(_radial_f, Verbose) if valid_function: _radial_f = self.get_func(_radial_f) @@ -1884,17 +1725,17 @@ def GetNpoints(self, Patch: Patch) -> tuple: """ poloidal_tag, radial_tag = Patch.get_tag() - np_tag = "np_" + poloidal_tag - nr_tag = "nr_" + radial_tag + np_tag = 'np_' + poloidal_tag + nr_tag = 'nr_' + radial_tag try: - np_cells = self.settings["grid_settings"]["grid_generation"][np_tag] + np_cells = self.settings['grid_settings']['grid_generation'][np_tag] except: - np_cells = self.settings["grid_settings"]["grid_generation"]["np_default"] + np_cells = self.settings['grid_settings']['grid_generation']['np_default'] try: - nr_cells = self.settings["grid_settings"]["grid_generation"][nr_tag] + nr_cells = self.settings['grid_settings']['grid_generation'][nr_tag] except: - nr_cells = self.settings["grid_settings"]["grid_generation"]["nr_default"] + nr_cells = self.settings['grid_settings']['grid_generation']['nr_default'] return (nr_cells, np_cells) @@ -1909,30 +1750,16 @@ def GetDistortionCorrectionSettings(self) -> dict: ------- A dictionary containing ``CorrectDistortion`` settings. """ - if ( - self.settings["grid_settings"]["grid_generation"].get( - "distortion_correction" - ) - is not None - ): - CD = self.settings["grid_settings"]["grid_generation"][ - "distortion_correction" - ] + if self.settings['grid_settings']['grid_generation'].get('distortion_correction') is not None: + CD = self.settings['grid_settings']['grid_generation']['distortion_correction'] else: CD = {} self.distortion_correction = CD return CD - def construct_grid( - self, - np_cells: int = 1, - nr_cells: int = 1, - Verbose: bool = False, - ShowVertices: bool = False, - RestartScratch: bool = False, - ListPatches: str = "all", - ) -> None: + def construct_grid(self, np_cells: int = 1, nr_cells: int = 1, Verbose: bool = False, + ShowVertices: bool = False, RestartScratch: bool = False, ListPatches: str = 'all') -> None: """ Construct a grid by refining a Patch map. @@ -1995,13 +1822,13 @@ def construct_grid( # Straighten up East and West segments of our patches, # Plot borders and fill patches. if Verbose: - print("Construct Grid") + print('Construct Grid') try: - visual = self.settings["DEBUG"]["visual"]["subgrid"] + visual = self.settings['DEBUG']['visual']['subgrid'] except: visual = False try: - verbose = self.settings["DEBUG"]["verbose"]["grid_generation"] + verbose = self.settings['DEBUG']['verbose']['grid_generation'] except: verbose = False @@ -2009,7 +1836,7 @@ def construct_grid( self.GetDistortionCorrectionSettings() - print(">>> Patches:", [k for k in self.patches.keys()]) + print('>>> Patches:', [k for k in self.patches.keys()]) if RestartScratch: self.CurrentListPatch = {} @@ -2018,37 +1845,24 @@ def construct_grid( if self.distortion_correction.get(name) is not None: patch.distortion_correction = self.distortion_correction.get(name) elif self.distortion_correction.get(patch.get_tag()) is not None: - patch.distortion_correction = self.distortion_correction.get( - patch.get_tag() - ) - elif self.distortion_correction.get("all") is not None: - patch.distortion_correction = self.distortion_correction.get("all") + patch.distortion_correction = self.distortion_correction.get(patch.get_tag()) + elif self.distortion_correction.get('all') is not None: + patch.distortion_correction = self.distortion_correction.get('all') else: - patch.distortion_correction = {"Active": False} - if (ListPatches == "all" and patch not in self.CurrentListPatch) or ( - ListPatches != "all" and name in ListPatches - ): + patch.distortion_correction = {'Active': False} + if (ListPatches == 'all' and patch not in self.CurrentListPatch) or (ListPatches != 'all' and name in ListPatches): self.SetPatchBoundaryPoints(patch, verbose) (nr_cells, np_cells) = self.GetNpoints(patch) (_radial_f, _poloidal_f) = self.GetFunctions(patch) - print(f">>> Making subgrid in patch {name}:") - print(f" np = {np_cells}, nr = {nr_cells}") - print(f" fp = {inspect.getsource(_poloidal_f)}") - print(f" fr = {inspect.getsource(_radial_f)}", end="") + print(f'>>> Making subgrid in patch {name}:') + print(f' np = {np_cells}, nr = {nr_cells}') + print(f' fp = {inspect.getsource(_poloidal_f)}') + print(f' fr = {inspect.getsource(_radial_f)}', end='') patch.RemoveDuplicatePoints() - patch.make_subgrid( - self, - np_cells, - nr_cells, - _poloidal_f=_poloidal_f, - _radial_f=_radial_f, - verbose=verbose, - visual=visual, - ShowVertices=ShowVertices, - ) + patch.make_subgrid(self, np_cells, nr_cells, _poloidal_f=_poloidal_f, _radial_f=_radial_f, verbose=verbose, visual=visual, ShowVertices=ShowVertices) self.CurrentListPatch[name] = patch - print(f" {name} subgrid complete.\n\n") + print(f' {name} subgrid complete.\n\n') self.AdjustGrid() def SetPatchBoundaryPoints(self, Patch: Patch, verbose: bool = False) -> None: @@ -2070,13 +1884,13 @@ def SetPatchBoundaryPoints(self, Patch: Patch, verbose: bool = False) -> None: Patch.TerminatesLoop = False if self.ConnexionMap.get(Patch.get_tag()) is not None: if verbose: - print("Find connexion map for patch {}".format(Patch.patch_name)) + print('Find connexion map for patch {}'.format(Patch.patch_name)) for v in self.ConnexionMap.get(Patch.get_tag()).items(): Boundary, AdjacentPatch = v Patch.BoundaryPoints[Boundary] = self.GetBoundaryPoints(AdjacentPatch) if verbose: - print("Find Boundaries points for {}".format(Patch.patch_name)) - if self.ConnexionMap.get(Patch.get_tag()).get("E") is not None: + print('Find Boundaries points for {}'.format(Patch.patch_name)) + if self.ConnexionMap.get(Patch.get_tag()).get('E') is not None: Patch.TerminatesLoop = True def GetBoundaryPoints(self, AdjacentPatchInfo: tuple) -> list: @@ -2099,13 +1913,13 @@ def GetBoundaryPoints(self, AdjacentPatchInfo: tuple) -> list: Boundary = AdjacentPatchInfo[1] for patch in self.patches.values(): if patch.get_tag() == PatchTag: - if Boundary == "S": + if Boundary == 'S': return patch.S_vertices - elif Boundary == "N": + elif Boundary == 'N': return patch.N_vertices - elif Boundary == "E": + elif Boundary == 'E': return patch.E_vertices - elif Boundary == "W": + elif Boundary == 'W': return patch.W_vertices return None @@ -2123,7 +1937,7 @@ def CheckPatches(self, verbose: bool = False) -> None: """ for name, patch in self.patches.items(): if patch.plate_patch: - print(" # Checking patch: ", name) + print(' # Checking patch: ', name) patch.CheckPatch(self) def SetupPatchMatrix(self) -> list: @@ -2159,235 +1973,30 @@ def SetupPatchMatrix(self) -> list: [[None], [None], [None], [None], [None], [None], [None], [None]] ] - elif self.config in ["SF45", "SF75", "SF105", "SF135"]: + elif self.config in ['SF45', 'SF75', 'SF105', 'SF135']: self.patch_matrix = [ - [ - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - ], - [ - [None], - p["A3"], - p["B3"], - p["C3"], - p["D3"], - p["E3"], - p["F3"], - p["G3"], - [None], - [None], - p["H3"], - p["I3"], - [None], - ], - [ - [None], - p["A2"], - p["B2"], - p["C2"], - p["D2"], - p["E2"], - p["F2"], - p["G2"], - [None], - [None], - p["H2"], - p["I2"], - [None], - ], - [ - [None], - p["A1"], - p["B1"], - p["C1"], - p["D1"], - p["E1"], - p["F1"], - p["G1"], - [None], - [None], - p["H1"], - p["I1"], - [None], - ], - [ - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - ], + [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]], + [[None], p['A3'], p['B3'], p['C3'], p['D3'], p['E3'], p['F3'], p['G3'], [None], [None], p['H3'], p['I3'], [None]], + [[None], p['A2'], p['B2'], p['C2'], p['D2'], p['E2'], p['F2'], p['G2'], [None], [None], p['H2'], p['I2'], [None]], + [[None], p['A1'], p['B1'], p['C1'], p['D1'], p['E1'], p['F1'], p['G1'], [None], [None], p['H1'], p['I1'], [None]], + [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]] ] - elif self.config in ["SF15", "SF165"]: + elif self.config in ['SF15', 'SF165']: self.patch_matrix = [ - [ - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - ], - [ - [None], - p["A3"], - p["B3"], - p["C3"], - p["D3"], - p["E3"], - p["F3"], - [None], - [None], - p["G3"], - p["H3"], - p["I3"], - [None], - ], - [ - [None], - p["A2"], - p["B2"], - p["C2"], - p["D2"], - p["E2"], - p["F2"], - [None], - [None], - p["G2"], - p["H2"], - p["I2"], - [None], - ], - [ - [None], - p["A1"], - p["B1"], - p["C1"], - p["D1"], - p["E1"], - p["F1"], - [None], - [None], - p["G1"], - p["H1"], - p["I1"], - [None], - ], - [ - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - ], + [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]], + [[None], p['A3'], p['B3'], p['C3'], p['D3'], p['E3'], p['F3'], [None], [None], p['G3'], p['H3'], p['I3'], [None]], + [[None], p['A2'], p['B2'], p['C2'], p['D2'], p['E2'], p['F2'], [None], [None], p['G2'], p['H2'], p['I2'], [None]], + [[None], p['A1'], p['B1'], p['C1'], p['D1'], p['E1'], p['F1'], [None], [None], p['G1'], p['H1'], p['I1'], [None]], + [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]] ] - elif self.config in ["UDN"]: + elif self.config in ['UDN']: self.patch_matrix = [ - [ - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - ], - [ - [None], - p["A3"], - p["B3"], - p["C3"], - p["D3"], - [None], - [None], - p["E3"], - p["F3"], - p["G3"], - p["H3"], - [None], - ], - [ - [None], - p["A2"], - p["B2"], - p["C2"], - p["D2"], - [None], - [None], - p["E2"], - p["F2"], - p["G2"], - p["H2"], - [None], - ], - [ - [None], - p["A1"], - p["B1"], - p["C1"], - p["D1"], - [None], - [None], - p["E1"], - p["F1"], - p["G1"], - p["H1"], - [None], - ], - [ - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - [None], - ], + [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]], + [[None], p['A3'], p['B3'], p['C3'], p['D3'], [None], [None], p['E3'], p['F3'], p['G3'], p['H3'], [None]], + [[None], p['A2'], p['B2'], p['C2'], p['D2'], [None], [None], p['E2'], p['F2'], p['G2'], p['H2'], [None]], + [[None], p['A1'], p['B1'], p['C1'], p['D1'], [None], [None], p['E1'], p['F1'], p['G1'], p['H1'], [None]], + [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]] ] return self.patch_matrix @@ -2415,76 +2024,40 @@ def _add_guardc(cell_map, ixlb, ixrb, nxpt=1, eps=1e-3): def set_guard(cell_map, ix, iy, eps, boundary): # Note: 'USN' and 'right' is really just 'LSN' and 'left' settings. # TODO: Edit the code to reflect this at some point so the next reader is not overwhelmed. - if boundary == "left": + if boundary == 'left': ixn = ix + 1 iyn = iy - cell_map[ix][iy][1] = cell_map[ixn][iyn][1] + eps * ( - cell_map[ixn][iyn][1] - cell_map[ixn][iyn][2] - ) + cell_map[ix][iy][1] = cell_map[ixn][iyn][1] + eps * (cell_map[ixn][iyn][1] - cell_map[ixn][iyn][2]) cell_map[ix][iy][2] = cell_map[ixn][iyn][1] - cell_map[ix][iy][3] = cell_map[ixn][iyn][3] + eps * ( - cell_map[ixn][iyn][3] - cell_map[ixn][iyn][4] - ) + cell_map[ix][iy][3] = cell_map[ixn][iyn][3] + eps * (cell_map[ixn][iyn][3] - cell_map[ixn][iyn][4]) cell_map[ix][iy][4] = cell_map[ixn][iyn][3] - cell_map[ix][iy][0] = 0.25 * ( - cell_map[ix][iy][1] - + cell_map[ix][iy][2] - + cell_map[ix][iy][3] - + cell_map[ix][iy][4] - ) - - elif boundary == "right": + cell_map[ix][iy][0] = 0.25 * (cell_map[ix][iy][1] + cell_map[ix][iy][2] + cell_map[ix][iy][3] + cell_map[ix][iy][4]) + + elif boundary == 'right': ixn = ix - 1 iyn = iy - cell_map[ix][iy][2] = cell_map[ixn][iyn][2] + eps * ( - cell_map[ixn][iyn][2] - cell_map[ixn][iyn][1] - ) + cell_map[ix][iy][2] = cell_map[ixn][iyn][2] + eps * (cell_map[ixn][iyn][2] - cell_map[ixn][iyn][1]) cell_map[ix][iy][1] = cell_map[ixn][iyn][2] - cell_map[ix][iy][4] = cell_map[ixn][iyn][4] + eps * ( - cell_map[ixn][iyn][4] - cell_map[ixn][iyn][3] - ) + cell_map[ix][iy][4] = cell_map[ixn][iyn][4] + eps * (cell_map[ixn][iyn][4] - cell_map[ixn][iyn][3]) cell_map[ix][iy][3] = cell_map[ixn][iyn][4] - cell_map[ix][iy][0] = 0.25 * ( - cell_map[ix][iy][1] - + cell_map[ix][iy][2] - + cell_map[ix][iy][3] - + cell_map[ix][iy][4] - ) - - elif boundary == "bottom": + cell_map[ix][iy][0] = 0.25 * (cell_map[ix][iy][1] + cell_map[ix][iy][2] + cell_map[ix][iy][3] + cell_map[ix][iy][4]) + + elif boundary == 'bottom': ixn = ix iyn = iy + 1 - cell_map[ix][iy][1] = cell_map[ixn][iyn][1] + eps * ( - cell_map[ixn][iyn][1] - cell_map[ixn][iyn][3] - ) + cell_map[ix][iy][1] = cell_map[ixn][iyn][1] + eps * (cell_map[ixn][iyn][1] - cell_map[ixn][iyn][3]) cell_map[ix][iy][3] = cell_map[ixn][iyn][1] - cell_map[ix][iy][2] = cell_map[ixn][iyn][2] + eps * ( - cell_map[ixn][iyn][2] - cell_map[ixn][iyn][4] - ) + cell_map[ix][iy][2] = cell_map[ixn][iyn][2] + eps * (cell_map[ixn][iyn][2] - cell_map[ixn][iyn][4]) cell_map[ix][iy][4] = cell_map[ixn][iyn][2] - cell_map[ix][iy][0] = 0.25 * ( - cell_map[ix][iy][1] - + cell_map[ix][iy][2] - + cell_map[ix][iy][3] - + cell_map[ix][iy][4] - ) - elif boundary == "top": + cell_map[ix][iy][0] = 0.25 * (cell_map[ix][iy][1] + cell_map[ix][iy][2] + cell_map[ix][iy][3] + cell_map[ix][iy][4]) + elif boundary == 'top': ixn = ix iyn = iy - 1 - cell_map[ix][iy][3] = cell_map[ixn][iyn][3] + eps * ( - cell_map[ixn][iyn][3] - cell_map[ixn][iyn][1] - ) + cell_map[ix][iy][3] = cell_map[ixn][iyn][3] + eps * (cell_map[ixn][iyn][3] - cell_map[ixn][iyn][1]) cell_map[ix][iy][1] = cell_map[ixn][iyn][3] - cell_map[ix][iy][4] = cell_map[ixn][iyn][4] + eps * ( - cell_map[ixn][iyn][4] - cell_map[ixn][iyn][2] - ) + cell_map[ix][iy][4] = cell_map[ixn][iyn][4] + eps * (cell_map[ixn][iyn][4] - cell_map[ixn][iyn][2]) cell_map[ix][iy][2] = cell_map[ixn][iyn][4] - cell_map[ix][iy][0] = 0.25 * ( - cell_map[ix][iy][1] - + cell_map[ix][iy][2] - + cell_map[ix][iy][3] - + cell_map[ix][iy][4] - ) + cell_map[ix][iy][0] = 0.25 * (cell_map[ix][iy][1] + cell_map[ix][iy][2] + cell_map[ix][iy][3] + cell_map[ix][iy][4]) return cell_map @@ -2493,15 +2066,15 @@ def set_guard(cell_map, ix, iy, eps, boundary): for iy in range(1, nr + 1): ix = ixlb - cell_map = set_guard(cell_map, ix, iy, eps, boundary="left") + cell_map = set_guard(cell_map, ix, iy, eps, boundary='left') ix = ixrb + 1 - cell_map = set_guard(cell_map, ix, iy, eps, boundary="right") + cell_map = set_guard(cell_map, ix, iy, eps, boundary='right') for ix in range(np + 2): iy = 0 - cell_map = set_guard(cell_map, ix, iy, eps, boundary="bottom") + cell_map = set_guard(cell_map, ix, iy, eps, boundary='bottom') iy = nr + 1 - cell_map = set_guard(cell_map, ix, iy, eps, boundary="top") + cell_map = set_guard(cell_map, ix, iy, eps, boundary='top') return cell_map @@ -2570,56 +2143,43 @@ def set_guard(cell_map, ix, iy, eps, boundary): elif (self.parent.settings['grid_settings']['num_xpt'] == 2) or ((self.parent.settings['grid_settings']['num_xpt'] == 1) and (self.parent.settings['grid_settings']["up_down_symmetry"] is True)): - if self.config in ["SF45", "SF75", "SF105", "SF135"]: + if self.config in ['SF45', 'SF75', 'SF105', 'SF135']: pindex1 = 8 pindex2 = 10 pindex3 = 12 - elif self.config in ["SF15", "SF165"]: + rindex = 4 + elif self.config in ['SF15', 'SF165']: pindex1 = 7 pindex2 = 9 pindex3 = 12 - elif self.config in ["UDN"]: + rindex = 4 + elif self.config in ['UDN']: pindex1 = 5 pindex2 = 7 pindex3 = 11 + rindex = 4 elif self.config in ["LSN", "USN"] and (self.parent.settings['grid_settings']["up_down_symmetry"] is True): pindex1 = 3 pindex2 = 5 pindex3 = 7 + rindex = 3 # Total number of poloidal indices in all subgrids. - np_total1 = ( - int(np.sum([patch.npol - 1 for patch in patch_matrix[1][1:pindex1]])) - + 2 - ) + np_total1 = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][1:pindex1]])) + 2 # Total number of radial indices in all subgrids. - # nr_total1 = int(np.sum([patch[1].nrad - 1 for patch in patch_matrix[1:4]])) + 2 - nr_total1 = ( - int(np.sum([patch[1].nrad - 1 for patch in patch_matrix[1:3]])) + 2 - ) + nr_total1 = int(np.sum([patch[1].nrad - 1 for patch in patch_matrix[1:rindex]])) + 2 # Total number of poloidal indices in all subgrids. - np_total2 = ( - int( - np.sum( - [patch.npol - 1 for patch in patch_matrix[1][pindex2:pindex3]] - ) - ) - + 2 - ) + np_total2 = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][pindex2:pindex3]])) + 2 # Total number of radial indices in all subgrids. - # nr_total2 = int(np.sum([patch[pindex2].nrad - 1 for patch in patch_matrix[1:4]])) + 2 - nr_total2 = ( - int(np.sum([patch[pindex2].nrad - 1 for patch in patch_matrix[1:3]])) - + 2 - ) + nr_total2 = int(np.sum([patch[pindex2].nrad - 1 for patch in patch_matrix[1:rindex]])) + 2 - rm1 = np.zeros((np_total1, nr_total1, 5), order="F") - zm1 = np.zeros((np_total1, nr_total1, 5), order="F") - rm2 = np.zeros((np_total2, nr_total2, 5), order="F") - zm2 = np.zeros((np_total2, nr_total2, 5), order="F") + rm1 = np.zeros((np_total1, nr_total1, 5), order='F') + zm1 = np.zeros((np_total1, nr_total1, 5), order='F') + rm2 = np.zeros((np_total2, nr_total2, 5), order='F') + zm2 = np.zeros((np_total2, nr_total2, 5), order='F') ixcell = 0 jycell = 0 @@ -2643,30 +2203,15 @@ def set_guard(cell_map, ix, iy, eps, boundary): # jyl - number of radial cells in the patch for jyl in range(len(local_patch.cell_grid)): - ixcell = ( - int( - np.sum( - [ - patch.npol - 1 - for patch in patch_matrix[1][1 : ixp + 1] - ] - ) - ) - - len(local_patch.cell_grid[0]) - + ixl - + 1 - ) + ixcell = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][1:ixp + 1]])) \ + - len(local_patch.cell_grid[0]) + ixl + 1 jycell = nr_sum - (local_patch.nrad - 1) + jyl + 1 ind = 0 - for coor in ["CENTER", "SW", "SE", "NW", "NE"]: - rm1[ixcell][jycell][ind] = ( - local_patch.cell_grid[jyl][ixl].vertices[coor].x - ) - zm1[ixcell][jycell][ind] = ( - local_patch.cell_grid[jyl][ixl].vertices[coor].y - ) + for coor in ['CENTER', 'SW', 'SE', 'NW', 'NE']: + rm1[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].x + zm1[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].y ind += 1 ixcell = 0 @@ -2690,32 +2235,15 @@ def set_guard(cell_map, ix, iy, eps, boundary): # jyl - number of radial cells in the patch for jyl in range(len(local_patch.cell_grid)): - ixcell = ( - int( - np.sum( - [ - patch.npol - 1 - for patch in patch_matrix[1][ - pindex2 : ixp + 1 - ] - ] - ) - ) - - len(local_patch.cell_grid[0]) - + ixl - + 1 - ) + ixcell = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][pindex2:ixp + 1]])) \ + - len(local_patch.cell_grid[0]) + ixl + 1 jycell = nr_sum - (local_patch.nrad - 1) + jyl + 1 ind = 0 - for coor in ["CENTER", "SW", "SE", "NW", "NE"]: - rm2[ixcell][jycell][ind] = ( - local_patch.cell_grid[jyl][ixl].vertices[coor].x - ) - zm2[ixcell][jycell][ind] = ( - local_patch.cell_grid[jyl][ixl].vertices[coor].y - ) + for coor in ['CENTER', 'SW', 'SE', 'NW', 'NE']: + rm2[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].x + zm2[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].y ind += 1 # Flip indices into gridue format. @@ -2743,7 +2271,7 @@ def set_guard(cell_map, ix, iy, eps, boundary): self.zm = np.concatenate((zm1, zm2)) try: - debug = self.yaml["DEBUG"]["visual"]["gridue"] + debug = self.yaml['DEBUG']['visual']['gridue'] except: debug = False From cd1f74754af73bbe781b7d2c8b4283584d69b6ee Mon Sep 17 00:00:00 2001 From: Dominic Power Date: Wed, 7 May 2025 16:52:28 -0400 Subject: [PATCH 03/10] Implemented updown symmetry from SF75 geometry --- INGRID/ingrid.py | 8 + INGRID/topologies/sf75.py | 75 +++-- INGRID/utils.py | 692 +++++++++++++++++++++++++------------- 3 files changed, 509 insertions(+), 266 deletions(-) diff --git a/INGRID/ingrid.py b/INGRID/ingrid.py index 686eb60..5fdd249 100644 --- a/INGRID/ingrid.py +++ b/INGRID/ingrid.py @@ -1432,7 +1432,15 @@ def ApplyUpDownSymmetry(self) -> None: del self.CurrentTopology.patches["C2"] del self.CurrentTopology.patches["D1"] del self.CurrentTopology.patches["D2"] + elif self.CurrentTopology.config == "SF75": + del self.CurrentTopology.patches["C1"] + del self.CurrentTopology.patches["C2"] + del self.CurrentTopology.patches["C3"] + del self.CurrentTopology.patches["D1"] + del self.CurrentTopology.patches["D2"] + del self.CurrentTopology.patches["D3"] else: + #TODO: Implement up/down symmetry for other geometries raise Exception("Up/down symmetry not yet implemented for config " + self.CurrentTopology.config) @_timer diff --git a/INGRID/topologies/sf75.py b/INGRID/topologies/sf75.py index eb350ce..da462ac 100644 --- a/INGRID/topologies/sf75.py +++ b/INGRID/topologies/sf75.py @@ -582,29 +582,58 @@ def set_gridue(self): ix_plate1 = 0 ix_cut1 = self.patches['A1'].npol - 1 - ix_cut2 = 0 - for alpha in ['A', 'B', 'C', 'D', 'E']: - ix_cut2 += self.patches[alpha + '1'].npol - 1 - - ix_plate2 = 0 - for alpha in ['A', 'B', 'C', 'D', 'E', 'F', 'G']: - ix_plate2 += self.patches[alpha + '3'].npol - 1 - - ix_plate3 = ix_plate2 + 2 - - ix_cut3 = 0 - for alpha in ['A', 'B', 'C', 'D', 'E', 'F']: - ix_cut3 += self.patches[alpha + '1'].npol - 1 - - ix_cut4 = 0 - for alpha in ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H']: - ix_cut4 += self.patches[alpha + '1'].npol - 1 - ix_cut4 += 2 - - ix_plate4 = 0 - for alpha in ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I']: - ix_plate4 += self.patches[alpha + '1'].npol - 1 - ix_plate4 += 2 + if self.settings["grid_settings"]["up_down_symmetry"] is True: + ix_cut2 = 0 + for alpha in ['A', 'B', 'E']: + ix_cut2 += self.patches[alpha + '1'].npol - 1 + ix_cut2 += 2 + + ix_plate2 = 0 + for alpha in ['A', 'B', 'E', 'F', 'G']: + ix_plate2 += self.patches[alpha + '3'].npol - 1 + ix_plate2 += 2 + + ix_plate3 = ix_plate2 + 2 + + ix_cut3 = 0 + for alpha in ['A', 'B', 'E', 'F']: + ix_cut3 += self.patches[alpha + '1'].npol - 1 + ix_cut3 += 2 + + ix_cut4 = 0 + for alpha in ['A', 'B', 'E', 'F', 'G', 'H']: + ix_cut4 += self.patches[alpha + '1'].npol - 1 + ix_cut4 += 4 + + ix_plate4 = 0 + for alpha in ['A', 'B', 'E', 'F', 'G', 'H', 'I']: + ix_plate4 += self.patches[alpha + '1'].npol - 1 + ix_plate4 += 4 + + elif self.settings["grid_settings"]["up_down_symmetry"] is False: + ix_cut2 = 0 + for alpha in ['A', 'B', 'C', 'D', 'E']: + ix_cut2 += self.patches[alpha + '1'].npol - 1 + + ix_plate2 = 0 + for alpha in ['A', 'B', 'C', 'D', 'E', 'F', 'G']: + ix_plate2 += self.patches[alpha + '3'].npol - 1 + + ix_plate3 = ix_plate2 + 2 + + ix_cut3 = 0 + for alpha in ['A', 'B', 'C', 'D', 'E', 'F']: + ix_cut3 += self.patches[alpha + '1'].npol - 1 + + ix_cut4 = 0 + for alpha in ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H']: + ix_cut4 += self.patches[alpha + '1'].npol - 1 + ix_cut4 += 2 + + ix_plate4 = 0 + for alpha in ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I']: + ix_plate4 += self.patches[alpha + '1'].npol - 1 + ix_plate4 += 2 psi = np.zeros((nxm + 2, nym + 2, 5), order='F') br = np.zeros((nxm + 2, nym + 2, 5), order='F') diff --git a/INGRID/utils.py b/INGRID/utils.py index 99754c1..e55834e 100644 --- a/INGRID/utils.py +++ b/INGRID/utils.py @@ -1958,7 +1958,6 @@ def SetupPatchMatrix(self) -> list: if self.config in ["LSN", "USN"]: if self.settings["grid_settings"]["up_down_symmetry"]: - #TODO: Implement up/down symmetry for other geometries self.patch_matrix = [ [[None], [None], [None], [None], [None], [None], [None], [None]], [[None], p["A2"], p["B2"], [None], [None], p["E2"], p["F2"], [None]], @@ -1974,13 +1973,22 @@ def SetupPatchMatrix(self) -> list: ] elif self.config in ['SF45', 'SF75', 'SF105', 'SF135']: - self.patch_matrix = [ - [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]], - [[None], p['A3'], p['B3'], p['C3'], p['D3'], p['E3'], p['F3'], p['G3'], [None], [None], p['H3'], p['I3'], [None]], - [[None], p['A2'], p['B2'], p['C2'], p['D2'], p['E2'], p['F2'], p['G2'], [None], [None], p['H2'], p['I2'], [None]], - [[None], p['A1'], p['B1'], p['C1'], p['D1'], p['E1'], p['F1'], p['G1'], [None], [None], p['H1'], p['I1'], [None]], - [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]] - ] + if self.settings["grid_settings"]["up_down_symmetry"]: + self.patch_matrix = [ + [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]], + [[None], p['A3'], p['B3'], [None], [None], p['E3'], p['F3'], p['G3'], [None], [None], p['H3'], p['I3'], [None]], + [[None], p['A2'], p['B2'], [None], [None], p['E2'], p['F2'], p['G2'], [None], [None], p['H2'], p['I2'], [None]], + [[None], p['A1'], p['B1'], [None], [None], p['E1'], p['F1'], p['G1'], [None], [None], p['H1'], p['I1'], [None]], + [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]] + ] + else: + self.patch_matrix = [ + [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]], + [[None], p['A3'], p['B3'], p['C3'], p['D3'], p['E3'], p['F3'], p['G3'], [None], [None], p['H3'], p['I3'], [None]], + [[None], p['A2'], p['B2'], p['C2'], p['D2'], p['E2'], p['F2'], p['G2'], [None], [None], p['H2'], p['I2'], [None]], + [[None], p['A1'], p['B1'], p['C1'], p['D1'], p['E1'], p['F1'], p['G1'], [None], [None], p['H1'], p['I1'], [None]], + [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]] + ] elif self.config in ['SF15', 'SF165']: self.patch_matrix = [ [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]], @@ -2016,264 +2024,462 @@ def concat_grid(self, guard_cell_eps: float = 1e-3) -> None: Determines the size of guard cell padding. """ - def _add_guardc(cell_map, ixlb, ixrb, nxpt=1, eps=1e-3): - """ - Add guard cells to the cell array. - """ - - def set_guard(cell_map, ix, iy, eps, boundary): - # Note: 'USN' and 'right' is really just 'LSN' and 'left' settings. - # TODO: Edit the code to reflect this at some point so the next reader is not overwhelmed. - if boundary == 'left': - ixn = ix + 1 - iyn = iy - cell_map[ix][iy][1] = cell_map[ixn][iyn][1] + eps * (cell_map[ixn][iyn][1] - cell_map[ixn][iyn][2]) - cell_map[ix][iy][2] = cell_map[ixn][iyn][1] - cell_map[ix][iy][3] = cell_map[ixn][iyn][3] + eps * (cell_map[ixn][iyn][3] - cell_map[ixn][iyn][4]) - cell_map[ix][iy][4] = cell_map[ixn][iyn][3] - cell_map[ix][iy][0] = 0.25 * (cell_map[ix][iy][1] + cell_map[ix][iy][2] + cell_map[ix][iy][3] + cell_map[ix][iy][4]) - - elif boundary == 'right': - ixn = ix - 1 - iyn = iy - cell_map[ix][iy][2] = cell_map[ixn][iyn][2] + eps * (cell_map[ixn][iyn][2] - cell_map[ixn][iyn][1]) - cell_map[ix][iy][1] = cell_map[ixn][iyn][2] - cell_map[ix][iy][4] = cell_map[ixn][iyn][4] + eps * (cell_map[ixn][iyn][4] - cell_map[ixn][iyn][3]) - cell_map[ix][iy][3] = cell_map[ixn][iyn][4] - cell_map[ix][iy][0] = 0.25 * (cell_map[ix][iy][1] + cell_map[ix][iy][2] + cell_map[ix][iy][3] + cell_map[ix][iy][4]) - - elif boundary == 'bottom': - ixn = ix - iyn = iy + 1 - cell_map[ix][iy][1] = cell_map[ixn][iyn][1] + eps * (cell_map[ixn][iyn][1] - cell_map[ixn][iyn][3]) - cell_map[ix][iy][3] = cell_map[ixn][iyn][1] - cell_map[ix][iy][2] = cell_map[ixn][iyn][2] + eps * (cell_map[ixn][iyn][2] - cell_map[ixn][iyn][4]) - cell_map[ix][iy][4] = cell_map[ixn][iyn][2] - cell_map[ix][iy][0] = 0.25 * (cell_map[ix][iy][1] + cell_map[ix][iy][2] + cell_map[ix][iy][3] + cell_map[ix][iy][4]) - elif boundary == 'top': - ixn = ix - iyn = iy - 1 - cell_map[ix][iy][3] = cell_map[ixn][iyn][3] + eps * (cell_map[ixn][iyn][3] - cell_map[ixn][iyn][1]) - cell_map[ix][iy][1] = cell_map[ixn][iyn][3] - cell_map[ix][iy][4] = cell_map[ixn][iyn][4] + eps * (cell_map[ixn][iyn][4] - cell_map[ixn][iyn][2]) - cell_map[ix][iy][2] = cell_map[ixn][iyn][4] - cell_map[ix][iy][0] = 0.25 * (cell_map[ix][iy][1] + cell_map[ix][iy][2] + cell_map[ix][iy][3] + cell_map[ix][iy][4]) - - return cell_map - - np = len(cell_map) - 2 - nr = len(cell_map[0]) - 2 - - for iy in range(1, nr + 1): - ix = ixlb - cell_map = set_guard(cell_map, ix, iy, eps, boundary='left') - ix = ixrb + 1 - cell_map = set_guard(cell_map, ix, iy, eps, boundary='right') - - for ix in range(np + 2): - iy = 0 - cell_map = set_guard(cell_map, ix, iy, eps, boundary='bottom') - iy = nr + 1 - cell_map = set_guard(cell_map, ix, iy, eps, boundary='top') - - return cell_map - patch_matrix = self.patch_matrix for patch in self.patches.values(): patch.npol = len(patch.cell_grid[0]) + 1 patch.nrad = len(patch.cell_grid) + 1 - if (self.parent.settings['grid_settings']['num_xpt'] == 1) and (self.parent.settings['grid_settings']["up_down_symmetry"] is False): + if self.parent.settings['grid_settings']["up_down_symmetry"] is True: + + if self.parent.settings['grid_settings']['num_xpt'] == 1: + self.rm, self.zm = self._concat_two_patch_regions(patch_matrix, pindex1 = 3, pindex2 = 5, pindex3 = 7, rindex = 3) + elif self.parent.settings['grid_settings']['num_xpt'] == 2: + if self.config in ['SF45', 'SF75', 'SF105', 'SF135']: + self.rm, self.zm = self._concat_three_patch_regions(patch_matrix, pindex1 = 3, pindex2 = 5, pindex3 = 8, pindex4 = 10, pindex5 = 12, rindex = 4) + elif self.config in ['SF15', 'SF165']: + # Note yet implemented + pass + elif self.config in ['UDN']: + # Not yet implemented + pass + + elif self.parent.settings['grid_settings']["up_down_symmetry"] is False: + if self.parent.settings['grid_settings']['num_xpt'] == 1: + self.rm, self.zm = self._concat_one_patch_region(patch_matrix) + elif self.parent.settings['grid_settings']['num_xpt'] == 2: + if self.config in ['SF45', 'SF75', 'SF105', 'SF135']: + self.rm, self.zm = self._concat_two_patch_regions(patch_matrix, pindex1 = 8, pindex2 = 10, pindex3 = 12, rindex = 4) + elif self.config in ['SF15', 'SF165']: + self.rm, self.zm = self._concat_two_patch_regions(patch_matrix, pindex1 = 7, pindex2 = 9, pindex3 = 12, rindex = 4) + elif self.config in ['UDN']: + self.rm, self.zm = self._concat_two_patch_regions(patch_matrix, pindex1 = 5, pindex2 = 7, pindex3 = 11, rindex = 4) + + def _add_guardc(self, cell_map, ixlb, ixrb, nxpt=1, eps=1e-3): + """ + Add guard cells to the cell array. + """ + + def set_guard(cell_map, ix, iy, eps, boundary): + # Note: 'USN' and 'right' is really just 'LSN' and 'left' settings. + # TODO: Edit the code to reflect this at some point so the next reader is not overwhelmed. + if boundary == 'left': + ixn = ix + 1 + iyn = iy + cell_map[ix][iy][1] = cell_map[ixn][iyn][1] + eps * (cell_map[ixn][iyn][1] - cell_map[ixn][iyn][2]) + cell_map[ix][iy][2] = cell_map[ixn][iyn][1] + cell_map[ix][iy][3] = cell_map[ixn][iyn][3] + eps * (cell_map[ixn][iyn][3] - cell_map[ixn][iyn][4]) + cell_map[ix][iy][4] = cell_map[ixn][iyn][3] + cell_map[ix][iy][0] = 0.25 * (cell_map[ix][iy][1] + cell_map[ix][iy][2] + cell_map[ix][iy][3] + cell_map[ix][iy][4]) + + elif boundary == 'right': + ixn = ix - 1 + iyn = iy + cell_map[ix][iy][2] = cell_map[ixn][iyn][2] + eps * (cell_map[ixn][iyn][2] - cell_map[ixn][iyn][1]) + cell_map[ix][iy][1] = cell_map[ixn][iyn][2] + cell_map[ix][iy][4] = cell_map[ixn][iyn][4] + eps * (cell_map[ixn][iyn][4] - cell_map[ixn][iyn][3]) + cell_map[ix][iy][3] = cell_map[ixn][iyn][4] + cell_map[ix][iy][0] = 0.25 * (cell_map[ix][iy][1] + cell_map[ix][iy][2] + cell_map[ix][iy][3] + cell_map[ix][iy][4]) + + elif boundary == 'bottom': + ixn = ix + iyn = iy + 1 + cell_map[ix][iy][1] = cell_map[ixn][iyn][1] + eps * (cell_map[ixn][iyn][1] - cell_map[ixn][iyn][3]) + cell_map[ix][iy][3] = cell_map[ixn][iyn][1] + cell_map[ix][iy][2] = cell_map[ixn][iyn][2] + eps * (cell_map[ixn][iyn][2] - cell_map[ixn][iyn][4]) + cell_map[ix][iy][4] = cell_map[ixn][iyn][2] + cell_map[ix][iy][0] = 0.25 * (cell_map[ix][iy][1] + cell_map[ix][iy][2] + cell_map[ix][iy][3] + cell_map[ix][iy][4]) + elif boundary == 'top': + ixn = ix + iyn = iy - 1 + cell_map[ix][iy][3] = cell_map[ixn][iyn][3] + eps * (cell_map[ixn][iyn][3] - cell_map[ixn][iyn][1]) + cell_map[ix][iy][1] = cell_map[ixn][iyn][3] + cell_map[ix][iy][4] = cell_map[ixn][iyn][4] + eps * (cell_map[ixn][iyn][4] - cell_map[ixn][iyn][2]) + cell_map[ix][iy][2] = cell_map[ixn][iyn][4] + cell_map[ix][iy][0] = 0.25 * (cell_map[ix][iy][1] + cell_map[ix][iy][2] + cell_map[ix][iy][3] + cell_map[ix][iy][4]) - np_total = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][1:-1]])) + 2 - nr_total = int(np.sum([patch[1].nrad - 1 for patch in patch_matrix[1:3]])) + 2 + return cell_map - rm = np.zeros((np_total, nr_total, 5), order='F') - zm = np.zeros((np_total, nr_total, 5), order='F') + np = len(cell_map) - 2 + nr = len(cell_map[0]) - 2 - ixcell = 0 - jycell = 0 + for iy in range(1, nr + 1): + ix = ixlb + cell_map = set_guard(cell_map, ix, iy, eps, boundary='left') + ix = ixrb + 1 + cell_map = set_guard(cell_map, ix, iy, eps, boundary='right') - # Iterate over all the patches in our SNL configuration (we exclude guard cells denoted by '[None]') - for ixp in range(1, 7): + for ix in range(np + 2): + iy = 0 + cell_map = set_guard(cell_map, ix, iy, eps, boundary='bottom') + iy = nr + 1 + cell_map = set_guard(cell_map, ix, iy, eps, boundary='top') - nr_sum = 0 - for jyp in range(1, 3): - # Point to the current patch we are operating on. - local_patch = patch_matrix[jyp][ixp] - nr_sum += local_patch.nrad - 1 + return cell_map + + def _concat_one_patch_region(self, patch_matrix: list): + """Concatenate the grid cells from all patches grouped as a single patch region - # Access the grid that is contained within this local_patch. - # ixl - number of poloidal cells in the patch. - for ixl in range(len(local_patch.cell_grid[0])): - # jyl - number of radial cells in the patch - for jyl in range(len(local_patch.cell_grid)): + :param patch_matrix: Matrix containing all patches including boundary regions + :return: rm, zm + """ + np_total = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][1:-1]])) + 2 + nr_total = int(np.sum([patch[1].nrad - 1 for patch in patch_matrix[1:3]])) + 2 - ixcell = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][1:ixp + 1]])) \ - - len(local_patch.cell_grid[0]) + ixl + 1 + rm = np.zeros((np_total, nr_total, 5), order='F') + zm = np.zeros((np_total, nr_total, 5), order='F') - jycell = nr_sum - (local_patch.nrad - 1) + jyl + 1 + ixcell = 0 + jycell = 0 - ind = 0 - for coor in ['CENTER', 'SW', 'SE', 'NW', 'NE']: - rm[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].x - zm[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].y - ind += 1 + # Iterate over all the patches in our SNL configuration (we exclude guard cells denoted by '[None]') + for ixp in range(1, 7): - # Flip indices into gridue format. - for i in range(len(rm)): - rm[i] = rm[i][::-1] - for i in range(len(zm)): - zm[i] = zm[i][::-1] + nr_sum = 0 + for jyp in range(1, 3): + # Point to the current patch we are operating on. + local_patch = patch_matrix[jyp][ixp] + nr_sum += local_patch.nrad - 1 - # Add guard cells to the concatenated grid. - ixrb = len(rm) - 2 - ixlb = 0 - self.rm = _add_guardc(rm, ixlb, ixrb) - self.zm = _add_guardc(zm, ixlb, ixrb) + # Access the grid that is contained within this local_patch. + # ixl - number of poloidal cells in the patch. + for ixl in range(len(local_patch.cell_grid[0])): + # jyl - number of radial cells in the patch + for jyl in range(len(local_patch.cell_grid)): - try: - debug = self.settings['DEBUG']['visual']['gridue'] - except: - debug = False + ixcell = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][1:ixp + 1]])) \ + - len(local_patch.cell_grid[0]) + ixl + 1 - if debug: - self._animate_grid() + jycell = nr_sum - (local_patch.nrad - 1) + jyl + 1 - elif (self.parent.settings['grid_settings']['num_xpt'] == 2) or ((self.parent.settings['grid_settings']['num_xpt'] == 1) and (self.parent.settings['grid_settings']["up_down_symmetry"] is True)): + ind = 0 + for coor in ['CENTER', 'SW', 'SE', 'NW', 'NE']: + rm[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].x + zm[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].y + ind += 1 - if self.config in ['SF45', 'SF75', 'SF105', 'SF135']: - pindex1 = 8 - pindex2 = 10 - pindex3 = 12 - rindex = 4 - elif self.config in ['SF15', 'SF165']: - pindex1 = 7 - pindex2 = 9 - pindex3 = 12 - rindex = 4 - elif self.config in ['UDN']: - pindex1 = 5 - pindex2 = 7 - pindex3 = 11 - rindex = 4 - elif self.config in ["LSN", "USN"] and (self.parent.settings['grid_settings']["up_down_symmetry"] is True): - pindex1 = 3 - pindex2 = 5 - pindex3 = 7 - rindex = 3 - - # Total number of poloidal indices in all subgrids. - np_total1 = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][1:pindex1]])) + 2 - - # Total number of radial indices in all subgrids. - nr_total1 = int(np.sum([patch[1].nrad - 1 for patch in patch_matrix[1:rindex]])) + 2 - - # Total number of poloidal indices in all subgrids. - np_total2 = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][pindex2:pindex3]])) + 2 - - # Total number of radial indices in all subgrids. - nr_total2 = int(np.sum([patch[pindex2].nrad - 1 for patch in patch_matrix[1:rindex]])) + 2 - - rm1 = np.zeros((np_total1, nr_total1, 5), order='F') - zm1 = np.zeros((np_total1, nr_total1, 5), order='F') - rm2 = np.zeros((np_total2, nr_total2, 5), order='F') - zm2 = np.zeros((np_total2, nr_total2, 5), order='F') - - ixcell = 0 - jycell = 0 - - # Iterate over all the patches in our DNL configuration (we exclude guard cells denoted by '[None]') - for ixp in range(1, pindex1): - - nr_sum = 0 - for jyp in range(1, 4): - # Point to the current patch we are operating on. - local_patch = patch_matrix[jyp][ixp] - - if local_patch == [None]: - continue - - nr_sum += local_patch.nrad - 1 - - # Access the grid that is contained within this local_patch. - # ixl - number of poloidal cells in the patch. - for ixl in range(len(local_patch.cell_grid[0])): - # jyl - number of radial cells in the patch - for jyl in range(len(local_patch.cell_grid)): - - ixcell = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][1:ixp + 1]])) \ - - len(local_patch.cell_grid[0]) + ixl + 1 - - jycell = nr_sum - (local_patch.nrad - 1) + jyl + 1 - - ind = 0 - for coor in ['CENTER', 'SW', 'SE', 'NW', 'NE']: - rm1[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].x - zm1[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].y - ind += 1 - - ixcell = 0 - jycell = 0 - - for ixp in range(pindex2, pindex3): - - nr_sum = 0 - for jyp in range(1, 4): - # Point to the current patch we are operating on. - local_patch = patch_matrix[jyp][ixp] - - if local_patch == [None]: - continue - - nr_sum += local_patch.nrad - 1 - - # Access the grid that is contained within this local_patch. - # ixl - number of poloidal cells in the patch. - for ixl in range(len(local_patch.cell_grid[0])): - # jyl - number of radial cells in the patch - for jyl in range(len(local_patch.cell_grid)): - - ixcell = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][pindex2:ixp + 1]])) \ - - len(local_patch.cell_grid[0]) + ixl + 1 - - jycell = nr_sum - (local_patch.nrad - 1) + jyl + 1 - - ind = 0 - for coor in ['CENTER', 'SW', 'SE', 'NW', 'NE']: - rm2[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].x - zm2[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].y - ind += 1 - - # Flip indices into gridue format. - for i in range(len(rm1)): - rm1[i] = rm1[i][::-1] - for i in range(len(zm1)): - zm1[i] = zm1[i][::-1] - for i in range(len(rm2)): - rm2[i] = rm2[i][::-1] - for i in range(len(zm2)): - zm2[i] = zm2[i][::-1] + # Flip indices into gridue format. + for i in range(len(rm)): + rm[i] = rm[i][::-1] + for i in range(len(zm)): + zm[i] = zm[i][::-1] - # Add guard cells to the concatenated grid. - ixrb1 = len(rm1) - 2 - ixlb1 = 0 - ixrb2 = len(rm2) - 2 - ixlb2 = 0 + # Add guard cells to the concatenated grid. + ixrb = len(rm) - 2 + ixlb = 0 + rm = self._add_guardc(rm, ixlb, ixrb) + zm = self._add_guardc(zm, ixlb, ixrb) - rm1 = _add_guardc(rm1, ixlb1, ixrb1) - zm1 = _add_guardc(zm1, ixlb1, ixrb1) - rm2 = _add_guardc(rm2, ixlb2, ixrb2) - zm2 = _add_guardc(zm2, ixlb2, ixrb2) + try: + debug = self.settings['DEBUG']['visual']['gridue'] + except: + debug = False - self.rm = np.concatenate((rm1, rm2)) - self.zm = np.concatenate((zm1, zm2)) + if debug: + self._animate_grid() - try: - debug = self.yaml['DEBUG']['visual']['gridue'] - except: - debug = False + return rm, zm + + def _concat_two_patch_regions(self, patch_matrix: list, pindex1: int, pindex2: int, pindex3: int, rindex: 3): + """Concatenate the grid cells from all patches grouped as two separate patch regions + + :param patch_matrix: Matrix containing all patches including boundary regions + :param pindex1: Poloidal patch index for the end of the first patch region + :param pindex2: Poloidal patch index for the start of the second patch region + :param pindex3: Poloidal patch index for the end of the second patch region + :param rindex: Number of radial patches + :return: rm, zm + """ + + # Total number of poloidal indices in all subgrids. + np_total1 = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][1:pindex1]])) + 2 + + # Total number of radial indices in all subgrids. + nr_total1 = int(np.sum([patch[1].nrad - 1 for patch in patch_matrix[1:rindex]])) + 2 + + # Total number of poloidal indices in all subgrids. + np_total2 = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][pindex2:pindex3]])) + 2 + + # Total number of radial indices in all subgrids. + nr_total2 = int(np.sum([patch[pindex2].nrad - 1 for patch in patch_matrix[1:rindex]])) + 2 + + rm1 = np.zeros((np_total1, nr_total1, 5), order='F') + zm1 = np.zeros((np_total1, nr_total1, 5), order='F') + rm2 = np.zeros((np_total2, nr_total2, 5), order='F') + zm2 = np.zeros((np_total2, nr_total2, 5), order='F') + + ixcell = 0 + jycell = 0 + + # Iterate over all the patches in our DNL configuration (we exclude guard cells denoted by '[None]') + for ixp in range(1, pindex1): + + nr_sum = 0 + for jyp in range(1, 4): + # Point to the current patch we are operating on. + local_patch = patch_matrix[jyp][ixp] + + if local_patch == [None]: + continue + + nr_sum += local_patch.nrad - 1 + + # Access the grid that is contained within this local_patch. + # ixl - number of poloidal cells in the patch. + for ixl in range(len(local_patch.cell_grid[0])): + # jyl - number of radial cells in the patch + for jyl in range(len(local_patch.cell_grid)): + + ixcell = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][1:ixp + 1]])) \ + - len(local_patch.cell_grid[0]) + ixl + 1 + + jycell = nr_sum - (local_patch.nrad - 1) + jyl + 1 + + ind = 0 + for coor in ['CENTER', 'SW', 'SE', 'NW', 'NE']: + rm1[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].x + zm1[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].y + ind += 1 + + ixcell = 0 + jycell = 0 + + for ixp in range(pindex2, pindex3): + + nr_sum = 0 + for jyp in range(1, 4): + # Point to the current patch we are operating on. + local_patch = patch_matrix[jyp][ixp] + + if local_patch == [None]: + continue + + nr_sum += local_patch.nrad - 1 + + # Access the grid that is contained within this local_patch. + # ixl - number of poloidal cells in the patch. + for ixl in range(len(local_patch.cell_grid[0])): + # jyl - number of radial cells in the patch + for jyl in range(len(local_patch.cell_grid)): + + ixcell = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][pindex2:ixp + 1]])) \ + - len(local_patch.cell_grid[0]) + ixl + 1 + + jycell = nr_sum - (local_patch.nrad - 1) + jyl + 1 + + ind = 0 + for coor in ['CENTER', 'SW', 'SE', 'NW', 'NE']: + rm2[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].x + zm2[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].y + ind += 1 + + # Flip indices into gridue format. + for i in range(len(rm1)): + rm1[i] = rm1[i][::-1] + for i in range(len(zm1)): + zm1[i] = zm1[i][::-1] + for i in range(len(rm2)): + rm2[i] = rm2[i][::-1] + for i in range(len(zm2)): + zm2[i] = zm2[i][::-1] + + # Add guard cells to the concatenated grid. + ixrb1 = len(rm1) - 2 + ixlb1 = 0 + ixrb2 = len(rm2) - 2 + ixlb2 = 0 + + rm1 = self._add_guardc(rm1, ixlb1, ixrb1) + zm1 = self._add_guardc(zm1, ixlb1, ixrb1) + rm2 = self._add_guardc(rm2, ixlb2, ixrb2) + zm2 = self._add_guardc(zm2, ixlb2, ixrb2) + + rm = np.concatenate((rm1, rm2)) + zm = np.concatenate((zm1, zm2)) + + try: + debug = self.yaml['DEBUG']['visual']['gridue'] + except: + debug = False + + if debug: + self._animate_grid() + + return rm, zm + + def _concat_three_patch_regions(self, patch_matrix: list, pindex1: int, pindex2: int, pindex3: int, pindex4: int, pindex5: int, rindex: int): + """Concatenate the grid cells from all patches grouped as three separate patch regions + + :param patch_matrix: Matrix containing all patches including boundary regions + :param pindex1: Poloidal patch index for the end of the first patch region + :param pindex2: Poloidal patch index for the start of the second patch region + :param pindex3: Poloidal patch index for the end of the second patch region + :param pindex3: Poloidal patch index for the start of the third patch region + :param pindex3: Poloidal patch index for the end of the third patch region + :param rindex: Number of radial patches + :return: rm, zm + """ + + # Total number of poloidal indices in all subgrids. + np_total1 = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][1:pindex1]])) + 2 + + # Total number of radial indices in all subgrids. + nr_total1 = int(np.sum([patch[1].nrad - 1 for patch in patch_matrix[1:rindex]])) + 2 + + # Total number of poloidal indices in all subgrids. + np_total2 = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][pindex2:pindex3]])) + 2 + + # Total number of radial indices in all subgrids. + nr_total2 = int(np.sum([patch[pindex2].nrad - 1 for patch in patch_matrix[1:rindex]])) + 2 + + # Total number of poloidal indices in all subgrids. + np_total3 = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][pindex4:pindex5]])) + 2 + + # Total number of radial indices in all subgrids. + nr_total3 = int(np.sum([patch[pindex4].nrad - 1 for patch in patch_matrix[1:rindex]])) + 2 + + rm1 = np.zeros((np_total1, nr_total1, 5), order='F') + zm1 = np.zeros((np_total1, nr_total1, 5), order='F') + rm2 = np.zeros((np_total2, nr_total2, 5), order='F') + zm2 = np.zeros((np_total2, nr_total2, 5), order='F') + rm3 = np.zeros((np_total3, nr_total3, 5), order='F') + zm3 = np.zeros((np_total3, nr_total3, 5), order='F') + + ixcell = 0 + jycell = 0 + + # Iterate over all the patches in our DNL configuration (we exclude guard cells denoted by '[None]') + for ixp in range(1, pindex1): + + nr_sum = 0 + for jyp in range(1, rindex): + # Point to the current patch we are operating on. + local_patch = patch_matrix[jyp][ixp] + + if local_patch == [None]: + continue + + nr_sum += local_patch.nrad - 1 + + # Access the grid that is contained within this local_patch. + # ixl - number of poloidal cells in the patch. + for ixl in range(len(local_patch.cell_grid[0])): + # jyl - number of radial cells in the patch + for jyl in range(len(local_patch.cell_grid)): + + ixcell = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][1:ixp + 1]])) \ + - len(local_patch.cell_grid[0]) + ixl + 1 + + jycell = nr_sum - (local_patch.nrad - 1) + jyl + 1 + + ind = 0 + for coor in ['CENTER', 'SW', 'SE', 'NW', 'NE']: + rm1[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].x + zm1[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].y + ind += 1 + + ixcell = 0 + jycell = 0 + + for ixp in range(pindex2, pindex3): + + nr_sum = 0 + for jyp in range(1, rindex): + # Point to the current patch we are operating on. + local_patch = patch_matrix[jyp][ixp] + + if local_patch == [None]: + continue + + nr_sum += local_patch.nrad - 1 + + # Access the grid that is contained within this local_patch. + # ixl - number of poloidal cells in the patch. + for ixl in range(len(local_patch.cell_grid[0])): + # jyl - number of radial cells in the patch + for jyl in range(len(local_patch.cell_grid)): + + ixcell = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][pindex2:ixp + 1]])) \ + - len(local_patch.cell_grid[0]) + ixl + 1 + + jycell = nr_sum - (local_patch.nrad - 1) + jyl + 1 + + ind = 0 + for coor in ['CENTER', 'SW', 'SE', 'NW', 'NE']: + rm2[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].x + zm2[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].y + ind += 1 + + ixcell = 0 + jycell = 0 + + for ixp in range(pindex4, pindex5): + + nr_sum = 0 + for jyp in range(1, rindex): + # Point to the current patch we are operating on. + local_patch = patch_matrix[jyp][ixp] + + if local_patch == [None]: + continue + + nr_sum += local_patch.nrad - 1 + + # Access the grid that is contained within this local_patch. + # ixl - number of poloidal cells in the patch. + for ixl in range(len(local_patch.cell_grid[0])): + # jyl - number of radial cells in the patch + for jyl in range(len(local_patch.cell_grid)): + + ixcell = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][pindex4:ixp + 1]])) \ + - len(local_patch.cell_grid[0]) + ixl + 1 + + jycell = nr_sum - (local_patch.nrad - 1) + jyl + 1 + + ind = 0 + for coor in ['CENTER', 'SW', 'SE', 'NW', 'NE']: + rm3[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].x + zm3[ixcell][jycell][ind] = local_patch.cell_grid[jyl][ixl].vertices[coor].y + ind += 1 + + # Flip indices into gridue format. + for i in range(len(rm1)): + rm1[i] = rm1[i][::-1] + for i in range(len(zm1)): + zm1[i] = zm1[i][::-1] + for i in range(len(rm2)): + rm2[i] = rm2[i][::-1] + for i in range(len(zm2)): + zm2[i] = zm2[i][::-1] + for i in range(len(rm3)): + rm3[i] = rm3[i][::-1] + for i in range(len(zm3)): + zm3[i] = zm3[i][::-1] + + # Add guard cells to the concatenated grid. + ixrb1 = len(rm1) - 2 + ixlb1 = 0 + ixrb2 = len(rm2) - 2 + ixlb2 = 0 + ixrb3 = len(rm3) - 2 + ixlb3 = 0 + + rm1 = self._add_guardc(rm1, ixlb1, ixrb1) + zm1 = self._add_guardc(zm1, ixlb1, ixrb1) + rm2 = self._add_guardc(rm2, ixlb2, ixrb2) + zm2 = self._add_guardc(zm2, ixlb2, ixrb2) + rm3 = self._add_guardc(rm3, ixlb3, ixrb3) + zm3 = self._add_guardc(zm3, ixlb3, ixrb3) + + rm = np.concatenate((rm1, rm2, rm3)) + zm = np.concatenate((zm1, zm2, zm3)) + + try: + debug = self.yaml['DEBUG']['visual']['gridue'] + except: + debug = False + + if debug: + self._animate_grid() - if debug: - self._animate_grid() + return rm, zm \ No newline at end of file From b31632d7d22e79ad08e6535e673448bc12ace2ef Mon Sep 17 00:00:00 2001 From: Dominic Power Date: Wed, 7 May 2025 17:18:07 -0400 Subject: [PATCH 04/10] Changed docstrings to default format --- INGRID/utils.py | 62 +++++++++++++++++++++++++++++++++++-------------- 1 file changed, 45 insertions(+), 17 deletions(-) diff --git a/INGRID/utils.py b/INGRID/utils.py index e55834e..0aa0708 100644 --- a/INGRID/utils.py +++ b/INGRID/utils.py @@ -2120,8 +2120,14 @@ def set_guard(cell_map, ix, iy, eps, boundary): def _concat_one_patch_region(self, patch_matrix: list): """Concatenate the grid cells from all patches grouped as a single patch region - :param patch_matrix: Matrix containing all patches including boundary regions - :return: rm, zm + Parameters + ---------- + patch_matrix: list + Matrix containing all patches including boundary regions + + Returns + ------- + rm, zm """ np_total = int(np.sum([patch.npol - 1 for patch in patch_matrix[1][1:-1]])) + 2 nr_total = int(np.sum([patch[1].nrad - 1 for patch in patch_matrix[1:3]])) + 2 @@ -2180,15 +2186,25 @@ def _concat_one_patch_region(self, patch_matrix: list): return rm, zm - def _concat_two_patch_regions(self, patch_matrix: list, pindex1: int, pindex2: int, pindex3: int, rindex: 3): + def _concat_two_patch_regions(self, patch_matrix: list, pindex1: int, pindex2: int, pindex3: int, rindex: int): """Concatenate the grid cells from all patches grouped as two separate patch regions - :param patch_matrix: Matrix containing all patches including boundary regions - :param pindex1: Poloidal patch index for the end of the first patch region - :param pindex2: Poloidal patch index for the start of the second patch region - :param pindex3: Poloidal patch index for the end of the second patch region - :param rindex: Number of radial patches - :return: rm, zm + Parameters + ---------- + patch_matrix: list + Matrix containing all patches including boundary regions + pindex1: int + Poloidal patch index for the end of the first patch region + pindex2: int + Poloidal patch index for the start of the second patch region + pindex3: int + Poloidal patch index for the end of the second patch region + rindex: int + Number of radial patches + + Returns + ------- + rm, zm """ # Total number of poloidal indices in all subgrids. @@ -2310,14 +2326,26 @@ def _concat_two_patch_regions(self, patch_matrix: list, pindex1: int, pindex2: i def _concat_three_patch_regions(self, patch_matrix: list, pindex1: int, pindex2: int, pindex3: int, pindex4: int, pindex5: int, rindex: int): """Concatenate the grid cells from all patches grouped as three separate patch regions - :param patch_matrix: Matrix containing all patches including boundary regions - :param pindex1: Poloidal patch index for the end of the first patch region - :param pindex2: Poloidal patch index for the start of the second patch region - :param pindex3: Poloidal patch index for the end of the second patch region - :param pindex3: Poloidal patch index for the start of the third patch region - :param pindex3: Poloidal patch index for the end of the third patch region - :param rindex: Number of radial patches - :return: rm, zm + Parameters + ---------- + patch_matrix: list + Matrix containing all patches including boundary regions + pindex1: int + Poloidal patch index for the end of the first patch region + pindex2: int + Poloidal patch index for the start of the second patch region + pindex3: int + Poloidal patch index for the end of the second patch region + pindex4: int + Poloidal patch index for the start of the third patch region + pindex5: int + Poloidal patch index for the end of the third patch region + rindex: int + Number of radial patches + + Returns + ------- + rm, zm """ # Total number of poloidal indices in all subgrids. From d40ae61bf6c35468e8acece35b0345c30011c377 Mon Sep 17 00:00:00 2001 From: Dominic Power Date: Mon, 4 Aug 2025 09:57:06 -0400 Subject: [PATCH 05/10] up/down symettry for SF15 --- INGRID/ingrid.py | 2 +- INGRID/topologies/sf15.py | 77 +++++++++++++++++++++++++++------------ INGRID/utils.py | 26 ++++++++----- 3 files changed, 71 insertions(+), 34 deletions(-) diff --git a/INGRID/ingrid.py b/INGRID/ingrid.py index 5fdd249..d19cbdd 100644 --- a/INGRID/ingrid.py +++ b/INGRID/ingrid.py @@ -1432,7 +1432,7 @@ def ApplyUpDownSymmetry(self) -> None: del self.CurrentTopology.patches["C2"] del self.CurrentTopology.patches["D1"] del self.CurrentTopology.patches["D2"] - elif self.CurrentTopology.config == "SF75": + elif self.CurrentTopology.config in ["SF15", "SF75"]: del self.CurrentTopology.patches["C1"] del self.CurrentTopology.patches["C2"] del self.CurrentTopology.patches["C3"] diff --git a/INGRID/topologies/sf15.py b/INGRID/topologies/sf15.py index 3672390..01a67e1 100644 --- a/INGRID/topologies/sf15.py +++ b/INGRID/topologies/sf15.py @@ -617,30 +617,59 @@ def set_gridue(self): ix_plate1 = 0 ix_cut1 = self.patches['A1'].npol - 1 - ix_cut2 = 0 - for alpha in ['A', 'B', 'C', 'D', 'E']: - ix_cut2 += self.patches[alpha + '1'].npol - 1 - - ix_plate2 = 0 - for alpha in ['A', 'B', 'C', 'D', 'E', 'F']: - ix_plate2 += self.patches[alpha + '3'].npol - 1 - - ix_plate3 = ix_plate2 + 2 - - ix_cut3 = 0 - for alpha in ['A', 'B', 'C', 'D', 'E', 'F', 'G']: - ix_cut3 += self.patches[alpha + '2'].npol - 1 - ix_cut3 += 2 - - ix_cut4 = 0 - for alpha in ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H']: - ix_cut4 += self.patches[alpha + '1'].npol - 1 - ix_cut4 += 2 - - ix_plate4 = 0 - for alpha in ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I']: - ix_plate4 += self.patches[alpha + '1'].npol - 1 - ix_plate4 += 2 + if self.settings["grid_settings"]["up_down_symmetry"] is True: + ix_cut2 = 0 + for alpha in ['A', 'B', 'E']: + ix_cut2 += self.patches[alpha + '1'].npol - 1 + ix_cut2 += 2 + + ix_plate2 = 0 + for alpha in ['A', 'B', 'E', 'F']: + ix_plate2 += self.patches[alpha + '3'].npol - 1 + ix_plate2 += 2 + + ix_plate3 = ix_plate2 + 2 + + ix_cut3 = 0 + for alpha in ['A', 'B', 'E', 'F', 'G']: + ix_cut3 += self.patches[alpha + '2'].npol - 1 + ix_cut3 += 4 + + ix_cut4 = 0 + for alpha in ['A', 'B', 'E', 'F', 'G', 'H']: + ix_cut4 += self.patches[alpha + '1'].npol - 1 + ix_cut4 += 4 + + ix_plate4 = 0 + for alpha in ['A', 'B', 'E', 'F', 'G', 'H', 'I']: + ix_plate4 += self.patches[alpha + '1'].npol - 1 + ix_plate4 += 4 + + elif self.settings["grid_settings"]["up_down_symmetry"] is False: + ix_cut2 = 0 + for alpha in ['A', 'B', 'C', 'D', 'E']: + ix_cut2 += self.patches[alpha + '1'].npol - 1 + + ix_plate2 = 0 + for alpha in ['A', 'B', 'C', 'D', 'E', 'F']: + ix_plate2 += self.patches[alpha + '3'].npol - 1 + + ix_plate3 = ix_plate2 + 2 + + ix_cut3 = 0 + for alpha in ['A', 'B', 'C', 'D', 'E', 'F', 'G']: + ix_cut3 += self.patches[alpha + '2'].npol - 1 + ix_cut3 += 2 + + ix_cut4 = 0 + for alpha in ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H']: + ix_cut4 += self.patches[alpha + '1'].npol - 1 + ix_cut4 += 2 + + ix_plate4 = 0 + for alpha in ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I']: + ix_plate4 += self.patches[alpha + '1'].npol - 1 + ix_plate4 += 2 psi = np.zeros((nxm + 2, nym + 2, 5), order='F') br = np.zeros((nxm + 2, nym + 2, 5), order='F') diff --git a/INGRID/utils.py b/INGRID/utils.py index 0aa0708..134dbc5 100644 --- a/INGRID/utils.py +++ b/INGRID/utils.py @@ -1990,13 +1990,22 @@ def SetupPatchMatrix(self) -> list: [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]] ] elif self.config in ['SF15', 'SF165']: - self.patch_matrix = [ - [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]], - [[None], p['A3'], p['B3'], p['C3'], p['D3'], p['E3'], p['F3'], [None], [None], p['G3'], p['H3'], p['I3'], [None]], - [[None], p['A2'], p['B2'], p['C2'], p['D2'], p['E2'], p['F2'], [None], [None], p['G2'], p['H2'], p['I2'], [None]], - [[None], p['A1'], p['B1'], p['C1'], p['D1'], p['E1'], p['F1'], [None], [None], p['G1'], p['H1'], p['I1'], [None]], - [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]] - ] + if self.settings["grid_settings"]["up_down_symmetry"]: + self.patch_matrix = [ + [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]], + [[None], p['A3'], p['B3'], [None], [None], p['E3'], p['F3'], [None], [None], p['G3'], p['H3'], p['I3'], [None]], + [[None], p['A2'], p['B2'], [None], [None], p['E2'], p['F2'], [None], [None], p['G2'], p['H2'], p['I2'], [None]], + [[None], p['A1'], p['B1'], [None], [None], p['E1'], p['F1'], [None], [None], p['G1'], p['H1'], p['I1'], [None]], + [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]] + ] + else: + self.patch_matrix = [ + [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]], + [[None], p['A3'], p['B3'], p['C3'], p['D3'], p['E3'], p['F3'], [None], [None], p['G3'], p['H3'], p['I3'], [None]], + [[None], p['A2'], p['B2'], p['C2'], p['D2'], p['E2'], p['F2'], [None], [None], p['G2'], p['H2'], p['I2'], [None]], + [[None], p['A1'], p['B1'], p['C1'], p['D1'], p['E1'], p['F1'], [None], [None], p['G1'], p['H1'], p['I1'], [None]], + [[None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None], [None]] + ] elif self.config in ['UDN']: self.patch_matrix = [ @@ -2038,8 +2047,7 @@ def concat_grid(self, guard_cell_eps: float = 1e-3) -> None: if self.config in ['SF45', 'SF75', 'SF105', 'SF135']: self.rm, self.zm = self._concat_three_patch_regions(patch_matrix, pindex1 = 3, pindex2 = 5, pindex3 = 8, pindex4 = 10, pindex5 = 12, rindex = 4) elif self.config in ['SF15', 'SF165']: - # Note yet implemented - pass + self.rm, self.zm = self._concat_three_patch_regions(patch_matrix, pindex1 = 3, pindex2 = 5, pindex3 = 7, pindex4 = 9, pindex5 = 12, rindex = 4) elif self.config in ['UDN']: # Not yet implemented pass From f7d5f98d77f16644216dd29416469d6ef78b0286 Mon Sep 17 00:00:00 2001 From: Dominic Power Date: Mon, 4 Aug 2025 09:58:41 -0400 Subject: [PATCH 06/10] Modified patch map debug tracing style --- INGRID/line_tracing.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/INGRID/line_tracing.py b/INGRID/line_tracing.py index 93209ec..6190bc1 100644 --- a/INGRID/line_tracing.py +++ b/INGRID/line_tracing.py @@ -652,7 +652,12 @@ def save_line(x, y, linecolor=color, marker='.-', markersize=1.5): if show_plot: if hasattr(self.grid, 'ax') is False: self.grid.plot_data() - self.grid.ax.plot(x, y, '.-', linewidth=2, color=color, markersize=1.5) + self.grid.ax.plot(x, y, '-', linewidth=0.5, color=color) + try: + self.cur_plot_point.set_ydata(y) + self.cur_plot_point.set_xdata(x) + except: + self.cur_plot_point, = self.grid.ax.plot(x, y, linestyle=None, marker="*", color="red", zorder=9999) plt.draw() plt.pause(np.finfo(float).eps) From 1e8eb62a1aa3c5f7ad0c1661feb244e2112d0ff7 Mon Sep 17 00:00:00 2001 From: Dominic Power Date: Mon, 25 Aug 2025 22:58:21 -0400 Subject: [PATCH 07/10] Incorporated psi extender for udsym grids into INGRID --- INGRID/ingrid.py | 2 +- INGRID/udsym_tools.py | 236 ++++++++++++++++++++++++++++++++++++++++++ INGRID/utils.py | 11 +- 3 files changed, 245 insertions(+), 4 deletions(-) create mode 100644 INGRID/udsym_tools.py diff --git a/INGRID/ingrid.py b/INGRID/ingrid.py index ec0facb..79a2f55 100644 --- a/INGRID/ingrid.py +++ b/INGRID/ingrid.py @@ -1371,7 +1371,7 @@ def StartSetup(self, **kwargs) -> None: + 'Must be <= 2).' raise ValueError(v_error_str) - self.LoadGEQDSK(self.settings['eqdsk']) + self.LoadGEQDSK(self.settings['eqdsk'], self.settings["grid_settings"]["up_down_symmetry"]) self.AutoRefineMagAxis() self.AutoRefineXPoint() if topology == 'DNL': diff --git a/INGRID/udsym_tools.py b/INGRID/udsym_tools.py new file mode 100644 index 0000000..fecedab --- /dev/null +++ b/INGRID/udsym_tools.py @@ -0,0 +1,236 @@ +import numpy as np +import matplotlib.pyplot as plt +from scipy import interpolate +from scipy.optimize import fsolve +import freeqdsk +import os + + +class PsinExtender: + """Class for extending psin from the lower half of a magnetic equilibrium to the upper half by matching field lines on the outer/inner side.""" + + def __init__(self, eqdsk_data: dict): + """Initialise + + :param filename: Full path of the magnetic equilibrium file + """ + self.eqdsk = eqdsk_data + self.r = self.eqdsk["r_grid"] + self.z = self.eqdsk["z_grid"] + self.psin = (self.eqdsk["psi"] - self.eqdsk["simagx"]) / (self.eqdsk["sibdry"] - self.eqdsk["simagx"]) + psin_shape = self.psin.shape + self.num_x = psin_shape[0] + self.num_y = psin_shape[1] + + def _get_monotonic_section(self, arr: np.ndarray) -> int: + """Find the monotonically increasing region of an input array which starts from the first element + + :param arr: Input array + :return: Index of the last element on the monotonically increasing region + """ + for i in range(1, len(arr)): + if arr[i] < arr[i - 1]: + max_idx = i + break + else: + max_idx = i + 1 + return max_idx + + def extend_psi(self): + """Extend psi from the middle of the domain into the upper half by tracing contours of psin from the inner to the outer side""" + r_modified = self.r.copy() + z_modified = self.z.copy() + psin_modified = self.psin.copy() + + num_x = psin_modified.shape[0] + num_y = psin_modified.shape[1] + + idx_magx, r_magx, z_magx = self._get_magnetic_axis( + r_modified, z_modified, psin_modified + ) + r_mid, psin_mid = self._get_mid_arrays(r_modified, psin_modified) + + self._interpfunc_r_psin_in, self._interpfunc_r_psin_out = ( + self._get_interp_funcs_r_psin(r_mid, psin_mid, idx_magx) + ) + self._interpfunc_psin_r_in, self._interpfunc_psin_r_out = ( + self._get_interp_funcs_psin_r(r_mid, psin_mid, idx_magx) + ) + + # Fill the upper half of psin array from psin_2D + for ix in range(num_x): + for jy in range(int(num_y / 2), num_y): + rval = r_modified[ix, jy] + zval = z_modified[ix, jy] + rhoval = np.sqrt((rval - r_magx) ** 2 + (zval - z_magx) ** 2) + thetaval = np.arctan2(zval - z_magx, rval - r_magx) + thetaval = np.pi - thetaval # -using zero at imid, pi at omid + psin_guess = np.clip( + self._psin_2d_simple(rhoval, thetaval), + np.min(psin_modified), + np.max(psin_modified), + ) + psin_modified[ix, jy] = psin_guess + # psinnval = self._psin_2d(rhoval, thetaval, psin_guess) + # psin_modified[ix, jy] = psinnval + + # Save the modified arrays + self.r_modified = r_modified + self.z_modified = z_modified + self.psin_modified = psin_modified + + self._update_eqdsk_data() + + def _update_eqdsk_data(self): + """Apply changes to R, Z and psi in the input eqdsk data""" + self.eqdsk["r_grid"] = self.r_modified + self.eqdsk["z_grid"] = self.z_modified + self.eqdsk["psi"] = self.psin_modified * (self.eqdsk["sibdry"] - self.eqdsk["simagx"]) + self.eqdsk["simagx"] + + def _get_mid_arrays(self, r: np.ndarray, psin: np.ndarray) -> np.ndarray: + """Get arrays of R and psin in the middle of the Z domain + + :param r: R coordinate array + :param psin: psin values array + :return: Arrays (R, psin) + """ + r_mid = r[:, int(self.num_y / 2)] + psin_mid = psin[:, int(self.num_y / 2)] + + return r_mid, psin_mid + + def _get_magnetic_axis( + self, r: np.ndarray, z: np.ndarray, psin: np.ndarray + ) -> tuple[int, float, float]: + """Find the location of the magnetic axis + + :param r: R coordinate array + :param z: Z coordinate array + :param psin: psin values array + :return: Location of the magnetic axis: (array_index, R, Z) + """ + psin_mid = psin[:, int(self.num_y / 2)] + idx_magx = np.argmin(psin_mid) + r_magx = r[idx_magx, int(self.num_y / 2)] + z_magx = z[idx_magx, int(self.num_y / 2)] + + return idx_magx, r_magx, z_magx + + def _get_outer_r_psin( + self, r_mid: np.ndarray, psin_mid: np.ndarray, idx_magx: int + ) -> tuple[np.ndarray, np.ndarray]: + """Get R and psin arrays from the magnetic axis to the outer midplane + + :param r_mid: Array of R values along the middle of the Z domain + :param psin_mid: Array of psin values along the middle of the Z domain + :param idx_magx: Index of the middle of the Z domain + :return: Arrays (R, psin) + """ + r_out = r_mid[idx_magx : self.num_x - 1] - r_mid[idx_magx] + psin_out = psin_mid[idx_magx : self.num_x - 1] - psin_mid[idx_magx] + return r_out, psin_out + + def _get_inner_r_psin( + self, r_mid: np.ndarray, psin_mid: np.ndarray, idx_magx: int + ) -> tuple[np.ndarray, np.ndarray]: + """Get R and psin arrays from the magnetic axis to the inner midplane (arrays will be reversed so psin is increasing) + + :param r_mid: Array of R values along the middle of the Z domain + :param psin_mid: Array of psin values along the middle of the Z domain + :param idx_magx: Index of the middle of the Z domain + :return: Arrays (R, psin) + """ + r_in = -(r_mid[0 : idx_magx + 1] - r_mid[idx_magx]) + psin_in = psin_mid[0 : idx_magx + 1] - psin_mid[idx_magx] + psin_in = psin_in[::-1] # reverse order, so it should (monotonically) increase + r_in = r_in[::-1] + return r_in, psin_in + + def _get_interp_funcs_r_psin( + self, r_mid: np.ndarray, psin_mid: np.ndarray, idx_magx: int + ) -> tuple[interpolate.CubicSpline, interpolate.CubicSpline]: + """Get functions which interpolate R as a function of psin, separately for the inner and outer sides + + :param r_mid: Array of R values along the middle of the Z domain + :param psin_mid: Array of psin values along the middle of the Z domain + :param idx_magx: Index of the middle of the Z domain + :return: Two interpolation functions: (inner side, outer side) + """ + + # Outer midplane part rho(psin) + r_out, psin_out = self._get_outer_r_psin(r_mid, psin_mid, idx_magx) + interpfunc_out = interpolate.CubicSpline(psin_out, r_out) + + # Inner midplane part rho(psin) + r_in, psin_in = self._get_inner_r_psin(r_mid, psin_mid, idx_magx) + max_idx = self._get_monotonic_section(psin_in) + interpfunc_in = interpolate.CubicSpline(psin_in[:max_idx], r_in[:max_idx]) + + return interpfunc_in, interpfunc_out + + def _get_interp_funcs_psin_r( + self, r_mid: np.ndarray, psin_mid: np.ndarray, idx_magx: int + ) -> tuple[interpolate.CubicSpline, interpolate.CubicSpline]: + """Get functions which interpolate psin as a function of R, separately for the inner and outer sides + + :param r_mid: Array of R values along the middle of the Z domain + :param psin_mid: Array of psin values along the middle of the Z domain + :param idx_magx: Index of the middle of the Z domain + :return: Two interpolation functions: (inner side, outer side) + """ + + # Outer midplane part rho(psin) + r_out, psin_out = self._get_outer_r_psin(r_mid, psin_mid, idx_magx) + interpfunc_out = interpolate.CubicSpline(r_out, psin_out) + + # Inner midplane part rho(psin) + r_in, psin_in = self._get_inner_r_psin(r_mid, psin_mid, idx_magx) + interpfunc_in = interpolate.CubicSpline(r_in, psin_in) + + return interpfunc_in, interpfunc_out + + def _rho_2d(self, psin: float, theta: float) -> float: + """Find coordinate rho at a given location (psin, theta) + + :param psin: psin + :param theta: Theta coordinate centred on magnetic axis + :return: Rho + """ + # Define rho(psin,theta) in the upper plane + rho_in = self._interpfunc_r_psin_in(psin) + rho_out = self._interpfunc_r_psin_out(psin) + res = rho_in + (rho_out - rho_in) * (theta / np.pi) + return res + + def _eqn(self, psin: float, *data: list[float]) -> float: + """Equation to be solved to find psin at a given location (rho, theta) + + :param psin: psin + :return: Equation result + """ + rho, theta = data + res = self._rho_2d(psin, theta) - rho + return res + + def _psin_2d(self, rho: float, theta: float, psin_0: float) -> float: + """Solve for psin the non-linear equation, rho_2d(psin, theta) = rho + + :param rho: Radial coordinate centred on magnetic axis + :param theta: Theta coordinated centred on magnetic axis + :return: Value of psin at coords (rho, theta) + """ + data = (rho, theta) + psin = fsolve(self._eqn, psin_0, args=data) + return psin + + def _psin_2d_simple(self, rhoval, thetaval): + """Solve for psin by interpolating between psin at the same rho either side of the midplane + + :param rho: Radial coordinate centred on magnetic axis + :param theta: Theta coordinated centred on magnetic axis + :return: Value of psin at coords (rho, theta) + """ + psin_in = self._interpfunc_psin_r_in(rhoval) + psin_out = self._interpfunc_psin_r_out(rhoval) + psinnval = (psin_in * (np.pi - thetaval) + psin_out * thetaval) / np.pi + return psinnval diff --git a/INGRID/utils.py b/INGRID/utils.py index 9b9a3d5..8fc6a06 100644 --- a/INGRID/utils.py +++ b/INGRID/utils.py @@ -25,7 +25,7 @@ from INGRID.interpol import EfitData from INGRID.line_tracing import LineTracing from INGRID.geometry import Point, Line, Patch, orientation_between - +from INGRID.udsym_tools import PsinExtender class IngridUtils: """ @@ -506,7 +506,7 @@ def ProcessPaths(self) -> None: ) continue - def LoadGEQDSK(self, geqdsk_path: str) -> None: + def LoadGEQDSK(self, geqdsk_path: str, up_down_symmetry: bool = False) -> None: """ Python class to read the psi data in from an ascii file. @@ -517,6 +517,11 @@ def LoadGEQDSK(self, geqdsk_path: str) -> None: geqdsk_data = geqdsk.read(f) if not isinstance(geqdsk_data,dict): geqdsk_data = geqdsk_data.__dict__ + + if up_down_symmetry: + # Chop off the top half of the domain and join lines of constant psi either side of the midplane + pe = PsinExtender(geqdsk_data) + pe.extend_psi() # # Extract quantities needed to initialize EfitData class @@ -749,7 +754,7 @@ def SetLimiter( # Empty list of coordinates falls back on using eqdsk limiter settings # else: - self.LoadGEQDSK(geqdsk_path=self.settings["eqdsk"]) + self.LoadGEQDSK(geqdsk_path=self.settings["eqdsk"], up_down_symmetry=self.settings["grid_settings"]["up_down_symmetry"]) self.geqdsk_data["rlim"] += rshift self.geqdsk_data["zlim"] += zshift From bceeb44e49b4400adea1f6d773ff58ca8ded1e95 Mon Sep 17 00:00:00 2001 From: Dominic Power Date: Mon, 25 Aug 2025 22:59:41 -0400 Subject: [PATCH 08/10] Fixed imports in udsym_tools module --- INGRID/udsym_tools.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/INGRID/udsym_tools.py b/INGRID/udsym_tools.py index fecedab..2755b68 100644 --- a/INGRID/udsym_tools.py +++ b/INGRID/udsym_tools.py @@ -1,9 +1,6 @@ import numpy as np -import matplotlib.pyplot as plt from scipy import interpolate from scipy.optimize import fsolve -import freeqdsk -import os class PsinExtender: From cc1d7a812516446260b489078348736ca69094f5 Mon Sep 17 00:00:00 2001 From: Dominic Power Date: Tue, 26 Aug 2025 11:36:17 -0400 Subject: [PATCH 09/10] Improved plotting when ud_down_symmetry option is True --- INGRID/ingrid.py | 23 ++++++++++++++++++----- INGRID/interpol.py | 11 ++++++++++- INGRID/udsym_tools.py | 2 +- INGRID/utils.py | 22 ++++++++++++++++++---- 4 files changed, 47 insertions(+), 11 deletions(-) diff --git a/INGRID/ingrid.py b/INGRID/ingrid.py index 79a2f55..7c8b90c 100644 --- a/INGRID/ingrid.py +++ b/INGRID/ingrid.py @@ -800,7 +800,9 @@ def PlotPsiUNorm(self) -> None: plt.draw() except: pass - self.PsiUNorm.plot_data(self.settings['grid_settings']['nlevs']) + self.PsiUNorm.plot_data(self.settings['grid_settings']['nlevs'], + up_down_symmetry=self.settings["grid_settings"]["up_down_symmetry"], + zmagx=self.magx[1] + self.settings["grid_settings"]["patch_generation"]["zmagx_shift"]) def PlotPsiNorm(self, view_mode: str = 'filled') -> None: """ @@ -813,7 +815,12 @@ def PlotPsiNorm(self, view_mode: str = 'filled') -> None: except: pass - self.PsiNorm.plot_data(nlevs=self.settings['grid_settings']['nlevs'], fig=self._PsiNormFig, ax=self.PsiNormAx, view_mode=view_mode) + self.PsiNorm.plot_data(nlevs=self.settings['grid_settings']['nlevs'], + fig=self._PsiNormFig, + ax=self.PsiNormAx, + view_mode=view_mode, + up_down_symmetry=self.settings["grid_settings"]["up_down_symmetry"], + zmagx=self.magx[1] + self.settings["grid_settings"]["patch_generation"]["zmagx_shift"]) def PlotPsiNormBounds(self) -> None: """ @@ -881,7 +888,7 @@ def PlotPsiNormMagReference(self, ax: object = None) -> None: x += self.settings['grid_settings']['patch_generation']['rmagx_shift'] y += self.settings['grid_settings']['patch_generation']['zmagx_shift'] self.RemovePlotPoint(label='magx', ax=ax) - ax.plot(x, y, '+', color='yellow', ms=15, linewidth=5, label='magx') + ax.plot(x, y, '+', color='yellow', ms=15, linewidth=5, label='magx', zorder=9999) (x, y) = self.xpt1 self.RemovePlotPoint(label='xpt1', ax=ax) @@ -1187,7 +1194,10 @@ def PlotPatches(self) -> None: pass self._PatchFig = plt.figure('INGRID: ' + self.CurrentTopology.config + ' Patches', figsize=(6, 10)) self.PatchAx = self._PatchFig.add_subplot(111) - self.CurrentTopology.patch_diagram(fig=self._PatchFig, ax=self.PatchAx) + self.CurrentTopology.patch_diagram(fig=self._PatchFig, + ax=self.PatchAx, + up_down_symmetry=self.settings["grid_settings"]["up_down_symmetry"], + zmagx=self.magx[1] + self.settings["grid_settings"]["patch_generation"]["zmagx_shift"]) self.PlotStrikeGeometry(ax=self.PatchAx) if self.settings['grid_settings']['patch_generation']['strike_pt_loc'] == 'target_plates': self.RemovePlotLine(label='limiter', ax=self.PatchAx) @@ -1202,7 +1212,10 @@ def PlotGrid(self) -> None: pass self._SubgridFig = plt.figure('INGRID: ' + self.CurrentTopology.config + ' Grid', figsize=(6, 10)) self._SubgridAx = self._SubgridFig.add_subplot(111) - self.CurrentTopology.grid_diagram(fig=self._SubgridFig, ax=self._SubgridAx) + self.CurrentTopology.grid_diagram(fig=self._SubgridFig, + ax=self._SubgridAx, + up_down_symmetry=self.settings["grid_settings"]["up_down_symmetry"], + zmagx=self.magx[1] + self.settings["grid_settings"]["patch_generation"]["zmagx_shift"]) def PlotSubgrid(self) -> None: """ diff --git a/INGRID/interpol.py b/INGRID/interpol.py index 3943dd1..0664e3a 100644 --- a/INGRID/interpol.py +++ b/INGRID/interpol.py @@ -4,6 +4,7 @@ """ import numpy as np import matplotlib.pyplot as plt +from matplotlib.patches import Rectangle from scipy.ndimage import zoom from scipy.interpolate import RectBivariateSpline as rbs @@ -240,7 +241,7 @@ def PlotLevel(self: object, level: float = 1.0, color: str = 'red', label: str = plt.plot([],[], label=label, color=color) def plot_data(self: object, nlevs: int = 30, interactive: bool = True, fig: object = None, - ax: object = None, view_mode: str = 'filled', refined: bool = True, refine_factor: int = 10): + ax: object = None, view_mode: str = 'filled', refined: bool = True, refine_factor: int = 10, up_down_symmetry: bool = False, zmagx: float = 0.0): """ Plot the EFIT data. @@ -264,6 +265,10 @@ def plot_data(self: object, nlevs: int = 30, interactive: bool = True, fig: obje Plot level with hi-resolution cubic spline representation refine_factor: int, optional Refinement factor for to be passed to SciPy zoom method + up_down_symmetry: bool, optional (defaults to False) + Up/down symmetry switch + zmagx: float, optional (defaults to 0.0) + Z position of magnetic axis (used in setting y-axis limits for the plot) """ lev = self.v.min() + (self.v.max() - self.v.min()) * np.arange(nlevs) / (nlevs - 1) @@ -289,6 +294,10 @@ def plot_data(self: object, nlevs: int = 30, interactive: bool = True, fig: obje self.ax.set_ylabel('Z') self.ax.set_xlim(self.rmin, self.rmax) self.ax.set_ylim(self.zmin, self.zmax) + if up_down_symmetry: + # Hide the upper half of the domain with a patch + self.ax.set_ylim(self.zmin, zmagx + 0.03*(self.zmax-self.zmin)) + ax.add_patch(Rectangle((self.rmin,zmagx), self.rmax - self.rmin, self.zmax - zmagx, facecolor="white", zorder=999)) if interactive: plt.ion() self.fig.show() diff --git a/INGRID/udsym_tools.py b/INGRID/udsym_tools.py index 2755b68..30f49a7 100644 --- a/INGRID/udsym_tools.py +++ b/INGRID/udsym_tools.py @@ -9,7 +9,7 @@ class PsinExtender: def __init__(self, eqdsk_data: dict): """Initialise - :param filename: Full path of the magnetic equilibrium file + :param eqdsk_data: Magnetic equilibrium data """ self.eqdsk = eqdsk_data self.r = self.eqdsk["r_grid"] diff --git a/INGRID/utils.py b/INGRID/utils.py index 8fc6a06..c34158d 100644 --- a/INGRID/utils.py +++ b/INGRID/utils.py @@ -1583,7 +1583,7 @@ def RefreshSettings(self): def OrderPatches(self): pass - def patch_diagram(self, fig: object = None, ax: object = None) -> None: + def patch_diagram(self, fig: object = None, ax: object = None, up_down_symmetry: bool = False, zmagx: float = 0.0) -> None: """ Generate the patch diagram for a given configuration. @@ -1594,6 +1594,10 @@ def patch_diagram(self, fig: object = None, ax: object = None) -> None: ax : object, optional Matplotlib axes to plot the Patch map on. + up_down_symmetry: bool, optional (defaults to False) + Up/down symmetry switch + zmagx: float, optional (defaults to 0.0) + Z position of magnetic axis (used in setting y-axis limits for the plot) """ @@ -1614,7 +1618,10 @@ def patch_diagram(self, fig: object = None, ax: object = None) -> None: f.subplots_adjust(bottom=0.2) a = ax if ax else f.subplots(1, 1) a.set_xlim([self.PsiUNorm.rmin, self.PsiUNorm.rmax]) - a.set_ylim([self.PsiUNorm.zmin, self.PsiUNorm.zmax]) + if up_down_symmetry: + a.set_ylim(self.PsiUNorm.zmin, zmagx + 0.03*(self.PsiUNorm.zmax-self.PsiUNorm.zmin)) + else: + a.set_ylim([self.PsiUNorm.zmin, self.PsiUNorm.zmax]) a.set_aspect("equal", adjustable="box") a.set_xlabel("R") @@ -1638,7 +1645,7 @@ def patch_diagram(self, fig: object = None, ax: object = None) -> None: ) f.show() - def grid_diagram(self, fig: object = None, ax: object = None) -> None: + def grid_diagram(self, fig: object = None, ax: object = None, up_down_symmetry: bool = False, zmagx: float = 0.0) -> None: """ Generates the grid diagram for a given configuration. @@ -1649,6 +1656,10 @@ def grid_diagram(self, fig: object = None, ax: object = None) -> None: ax : object, optional Matplotlib axes to plot the grid diagram on. + up_down_symmetry: bool, optional (defaults to False) + Up/down symmetry switch + zmagx: float, optional (defaults to 0.0) + Z position of magnetic axis (used in setting y-axis limits for the plot) """ if fig is None: @@ -1660,7 +1671,10 @@ def grid_diagram(self, fig: object = None, ax: object = None) -> None: plt.figure(fig.number) fig.subplots_adjust(bottom=0.2) ax.set_xlim(self.PsiUNorm.rmin, self.PsiUNorm.rmax) - ax.set_ylim(self.PsiUNorm.zmin, self.PsiUNorm.zmax) + if up_down_symmetry: + ax.set_ylim(self.PsiUNorm.zmin, zmagx + 0.03*(self.PsiUNorm.zmax-self.PsiUNorm.zmin)) + else: + ax.set_ylim(self.PsiUNorm.zmin, self.PsiUNorm.zmax) ax.set_aspect("equal", adjustable="box") ax.set_xlabel("R") From 2c8838bcb8a13665e5df8fb2c25dd846b0843d1f Mon Sep 17 00:00:00 2001 From: Dominic Power Date: Tue, 26 Aug 2025 11:49:44 -0400 Subject: [PATCH 10/10] Added example files for up_down_symmetry option in SNL and UDN geometries --- example_files/SNL/DIIID_SNL_bot.yml | 120 ++++++++++++++++++++++++++++ example_files/UDN/DIIID_UDN_bot.yml | 79 ++++++++++++++++++ 2 files changed, 199 insertions(+) create mode 100644 example_files/SNL/DIIID_SNL_bot.yml create mode 100644 example_files/UDN/DIIID_UDN_bot.yml diff --git a/example_files/SNL/DIIID_SNL_bot.yml b/example_files/SNL/DIIID_SNL_bot.yml new file mode 100644 index 0000000..e0947f5 --- /dev/null +++ b/example_files/SNL/DIIID_SNL_bot.yml @@ -0,0 +1,120 @@ +# --------------------------------------------------- +# User data directories +# --------------------------------------------------- +dir_settings: + eqdsk: ../data/SNL/DIII-D/ # dir containing eqdsk + limiter: . # dir containing limiter + patch_data: ../data/SNL/DIII-D/ # dir containing patch data + target_plates: ../data/SNL/DIII-D/ # dir containing target plates + +# --------------------------------------------------- +# eqdsk file name +# --------------------------------------------------- +eqdsk: neqdsk + +# --------------------------------------------------- +# General grid settings +# --------------------------------------------------- +grid_settings: + # ---------------------------------------------------------------------------- + # Settings for grid generation (num cells, transforms, distortion_correction) + # ---------------------------------------------------------------------------- + grid_generation: + distortion_correction: + all: + active: True # true, 1 also valid. + resolution: 1000 + theta_max: 140.0 + theta_min: 80.0 + F2: + active: False + np_A: 9 + np_F: 9 + np_B: 18 + np_E: 18 + + nr_2: 6 + + np_default: 3 + nr_default: 3 + poloidal_f_default: x, x + radial_f_default: x, x + # --------------------------------------------------- + # guard cell size + # --------------------------------------------------- + guard_cell_eps: 0.001 + # --------------------------------------------------- + # num levels in efit plot + # --------------------------------------------------- + nlevs: 30 + # --------------------------------------------------- + # make the grid up/down symmetric + # --------------------------------------------------- + up_down_symmetry: true + # --------------------------------------------------- + # num xpts + # --------------------------------------------------- + num_xpt: 1 + patch_generation: + strike_pt_loc: target_plates # 'limiter' or 'target_plates' + rmagx_shift: 0.0 + zmagx_shift: 0.0 + # --------------------------------------------------- + # Psi levels + # --------------------------------------------------- + psi_1: 1.066 + psi_core: 0.95 + psi_pf_1: 0.975 + # --------------------------------------------------- + # magx coordinates + # --------------------------------------------------- + rmagx: 1.75785604 + zmagx: -0.0292478683 + # --------------------------------------------------- + # xpt coordinates + # --------------------------------------------------- + rxpt: 1.300094032687 + zxpt: -1.133159375302 + # --------------------------------------------------- + # Filled contours vs contour lines + # --------------------------------------------------- + view_mode: filled + +# --------------------------------------------------- +# Saved patch settings +# --------------------------------------------------- +patch_data: + file: LSN_patches_1597099640.npy + preferences: + new_file: true + new_fname: LSN_patches_1597099640.npy + use_file: false + +# --------------------------------------------------- +# Integrator +# --------------------------------------------------- +integrator_settings: + dt: 0.01 + eps: 5.0e-06 + first_step: 5.0e-05 + max_step: 0.064 + step_ratio: 0.02 + tol: 0.005 + +# --------------------------------------------------- +# Limiter settings +# --------------------------------------------------- +limiter: + file: '' + use_efit_bounds: false + +# --------------------------------------------------- +# target plate settings +# --------------------------------------------------- +target_plates: + plate_E1: + file: d3d_otp.txt + zshift: -1.6 + plate_W1: + file: d3d_itp.txt + zshift: -1.6 diff --git a/example_files/UDN/DIIID_UDN_bot.yml b/example_files/UDN/DIIID_UDN_bot.yml new file mode 100644 index 0000000..5304883 --- /dev/null +++ b/example_files/UDN/DIIID_UDN_bot.yml @@ -0,0 +1,79 @@ +DEBUG: + verbose: + grid_generation: 0 + patch_generation: 0 + target_plates: 0 + visual: + find_NSEW: 0 + gridue: 0 + patch_map: 0 + subgrid: 0 +eqdsk: ../data/UDN/DIII-D/neqdsk +grid_settings: + grid_generation: + distortion_correction: + all: + active: false + resolution: 1000 + theta_max: 120.0 + theta_min: 80.0 + np_default: 2 + nr_default: 2 + poloidal_f_default: x, x + radial_f_default: x, x + guard_cell_eps: 0.001 + nlevs: 30 + num_xpt: 1 + up_down_symmetry: true + patch_generation: + xpt1_E_tilt: 0.0 + xpt1_W_tilt: 0.65 + rmagx_shift: 0.0 + xpt2_E_tilt: 0.0 + xpt2_W_tilt: 0.0 + strike_pt_loc: limiter + magx_tilt_1: 0.0 + magx_tilt_2: 0.0 + use_xpt1_E: false + use_xpt1_W: false + use_xpt2_E: false + use_xpt2_W: false + zmagx_shift: 0.0 + psi_1: 1.127 + psi_2: 1.131758078227 + psi_core: 0.975770930045 + psi_pf_1: 0.987326087084 + psi_pf_2: 1.1167326087084 + rmagx: 1.75785604 + rxpt: 1.300094032688 + rxpt2: 1.217505027591 + view_mode: filled + zmagx: -0.0292478683 + zxpt: -1.133159375304 + zxpt2: 1.219825715749 +integrator_settings: + dt: 0.01 + eps: 5.0e-05 + first_step: 5.0e-05 + max_step: 0.064 + step_ratio: 0.02 + tol: 0.005 +limiter: + use_efit_bounds: True + efit_buffer_r: 0.2 + efit_buffer_z: 0.2 +target_plates: + plate_E1: + file: ../data/UDN/DIII-D/d3d_otp.txt + zshift: -1.6 + plate_E2: + file: ../data/UDN/SAS1/itp4.txt + zshift: 0.0 + rshift: 0.05 + plate_W1: + file: ../data/UDN/DIII-D/d3d_itp.txt + zshift: -1.6 + plate_W2: + file: ../data/UDN/SAS1/SAS_odt.txt + zshift: 0.2 + rshift: -0.145