diff --git a/freegs4e/coil.py b/freegs4e/coil.py index 12edea2..5d4ae4e 100644 --- a/freegs4e/coil.py +++ b/freegs4e/coil.py @@ -132,6 +132,26 @@ def createPsiGreens(self, R, Z): """ return self.controlPsi(R, Z) + def createPsiGreensVec(self, R, Z): + """ + Calculate the Greens function at every point, and return + array. This will be passed back to evaluate Psi in + calcPsiFromGreens() + """ + return self.controlPsi(R, Z) + + def createBrGreensVec(self, R, Z): + """ + Calculate Br Greens functions + """ + return self.controlBr(R, Z) + + def createBzGreensVec(self, R, Z): + """ + Calculate Bz Greens functions + """ + return self.controlBz(R, Z) + def calcPsiFromGreens(self, pgreen): """ Calculate plasma psi from Greens functions and current diff --git a/freegs4e/critical.py b/freegs4e/critical.py index 037ccd8..c840617 100644 --- a/freegs4e/critical.py +++ b/freegs4e/critical.py @@ -404,7 +404,7 @@ def scan_for_crit(R, Z, psi): if det != 0: delta_R = -(fR * fZZ - 0.5 * fRZ * fZ) / det delta_Z = -(fZ * fRR - 0.5 * fRZ * fR) / det - if np.abs(delta_R) <= dR and np.abs(delta_Z) <= dZ: + if np.abs(delta_R) < 1.5 * dR and np.abs(delta_Z) < 1.5 * dZ: est_psi = psi[i, j] + 0.5 * ( fR * delta_R + fZ * delta_Z ) # + 0.5*(fRR*delta_R**2 + fZZ*delta_Z**2 + fRZ*delta_R*delta_Z) @@ -819,9 +819,12 @@ def inside_mask( use_geom=True, ): """ - This function calls inside_mask_ to find plasma region inside the separatrix (with the option of calling geom_inside_mask too). + geom_inside_mask applies an additional geometrical contraint + aimed at resolving cases of 'flooding' of the core mask through the primary Xpoint. + It excludes regions based on perpendicular to segment + from O-point to primary X-point. Parameters ---------- diff --git a/freegs4e/equilibrium.py b/freegs4e/equilibrium.py index 73d03c7..da9cfab 100644 --- a/freegs4e/equilibrium.py +++ b/freegs4e/equilibrium.py @@ -134,6 +134,7 @@ def __init__( # generate Greens function mappings (used # in self.psi() to speed up calculations) self._pgreen = tokamak.createPsiGreens(self.R, self.Z) + self._vgreen = tokamak.createPsiGreensVec(self.R, self.Z) # self._updatePlasmaPsi(psi) # Needs to be after _pgreen # assign plasma current diff --git a/freegs4e/jtor.py b/freegs4e/jtor.py index 97c2986..96dff25 100644 --- a/freegs4e/jtor.py +++ b/freegs4e/jtor.py @@ -163,12 +163,13 @@ def Jtor_part1(self, R, Z, psi, psi_bndry=None, mask_outside_limiter=None): R, Z, psi, self.mask_inside_limiter, self.Ip ) - # find core plasma mask (using psi_bndry) + # find core plasma mask (using user-defined psi_bndry) if psi_bndry is not None: diverted_core_mask = critical.inside_mask( R, Z, psi, opt, xpt, mask_outside_limiter, psi_bndry ) elif len(xpt) > 0: + # find core plasma mask using psi_bndry from xpt psi_bndry = xpt[0][2] self.psi_axis = opt[0][2] # check correct sorting between psi_axis and psi_bndry @@ -179,6 +180,37 @@ def Jtor_part1(self, R, Z, psi, psi_bndry=None, mask_outside_limiter=None): diverted_core_mask = critical.inside_mask( R, Z, psi, opt, xpt, mask_outside_limiter, psi_bndry ) + # check of any abrupt change of size in the diverted core mask + if hasattr(self, "diverted_core_mask"): + if self.diverted_core_mask is not None: + previous_core_size = np.sum(self.diverted_core_mask) + # check size change + check = ( + np.sum(diverted_core_mask) / previous_core_size < 0.5 + ) + check *= len(xpt) > 1 + if check: + # try using second xpt as primary xpt + alt_diverted_core_mask = critical.inside_mask( + R, + Z, + psi, + opt, + xpt[1:], + mask_outside_limiter, + xpt[1, 2], + ) + # check the alternative Xpoint gives rise to a valid core + edge_pixels = np.sum( + self.edge_mask * alt_diverted_core_mask + ) + if edge_pixels == 0: + print( + "Discarding 'primary' Xpoint! Please check final result" + ) + xpt = xpt[1:] + psi_bndry = xpt[1, 2] + diverted_core_mask = alt_diverted_core_mask.copy() else: # No X-points psi_bndry = psi[0, 0] @@ -547,7 +579,11 @@ def pprime(self, pn): """ shape = (1.0 - np.clip(pn, 0.0, 1.0) ** self.alpha_m) ** self.alpha_n - + if hasattr(self, "L") is False: + self.L = 1 + print( + "This is using self.L=1, which is likely not appropriate. Please calculate Jtor first to ensure the correct normalization." + ) return self.L * self.Beta0 / self.Raxis * shape def ffprime(self, pn): @@ -567,7 +603,11 @@ def ffprime(self, pn): """ shape = (1.0 - np.clip(pn, 0.0, 1.0) ** self.alpha_m) ** self.alpha_n - + if hasattr(self, "L") is False: + self.L = 1 + print( + "This is using self.L=1, which is likely not appropriate. Please calculate Jtor first to ensure the correct normalization." + ) return mu0 * self.L * (1 - self.Beta0) * self.Raxis * shape def fvac(self): @@ -713,7 +753,11 @@ def pprime(self, pn): """ shape = (1.0 - np.clip(pn, 0.0, 1.0) ** self.alpha_m) ** self.alpha_n - + if hasattr(self, "L") is False: + self.L = 1 + print( + "This is using self.L=1, which is likely not appropriate. Please calculate Jtor first to ensure the correct normalization." + ) return self.L * self.Beta0 / self.Raxis * shape def ffprime(self, pn): @@ -733,7 +777,11 @@ def ffprime(self, pn): """ shape = (1.0 - np.clip(pn, 0.0, 1.0) ** self.alpha_m) ** self.alpha_n - + if hasattr(self, "L") is False: + self.L = 1 + print( + "This is using self.L=1, which is likely not appropriate. Please calculate Jtor first to ensure the correct normalization." + ) return mu0 * self.L * (1 - self.Beta0) * self.Raxis * shape def fvac(self): @@ -786,9 +834,9 @@ def __init__( beta : list or np.array Polynomial coefficients for ffprime. alpha_logic : bool - If True, include the n_{P+1} term (see Jtor2). + If True, include the n_{P+1} term to ensure pprime(1)=0 (see Jtor2). beta_logic : bool - If True, include the n_{F+1} term (see Jtor2). + If True, include the n_{F+1} term to ensure ffprime(1)=0 (see Jtor2). Raxis : float Radial scaling parameter (non-negative). Ip_logic : bool @@ -833,7 +881,17 @@ def initialize_profile(self): self.beta = np.concatenate((self.beta, [-np.sum(self.beta)])) self.beta_exp = np.arange(0, len(self.beta)) - def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): + def Jtor_part2( + self, + R, + Z, + psi, + psi_axis, + psi_bndry, + mask, + torefine=False, + refineR=None, + ): """ Second part of the calculation that will use the explicit parameterisation of the chosen profile function to calculate Jtor. @@ -864,6 +922,7 @@ def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): mask : np.ndarray Mask of points inside the last closed flux surface. + Returns ------- np.array @@ -879,6 +938,9 @@ def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): dR = R[1, 0] - R[0, 0] dZ = Z[0, 1] - Z[0, 0] + if torefine: + R = 1.0 * refineR + # calculate normalised psi psi_norm = (psi - psi_axis) / (psi_bndry - psi_axis) psi_norm = np.clip(psi_norm, 0.0, 1.0) @@ -904,7 +966,14 @@ def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): # sum together Jtor = pprime_term + ffprime_term - # if there is a masking function, use it + # put to zero all current outside the LCFS + Jtor *= psi > psi_bndry + + Jtor *= self.Ip * Jtor > 0 + + if torefine: + return Jtor + if mask is not None: Jtor *= mask @@ -953,6 +1022,11 @@ def pprime(self, pn): list(np.shape(self.alpha)) + [1] * len(shape_pn) ) shape = np.sum(shape, axis=0) + if hasattr(self, "L") is False: + self.L = 1 + print( + "This is using self.L=1, which is likely not appropriate. Please calculate Jtor first to ensure the correct normalization." + ) return self.L * shape / self.Raxis def ffprime(self, pn): @@ -981,6 +1055,11 @@ def ffprime(self, pn): list(np.shape(self.beta)) + [1] * len(shape_pn) ) shape = np.sum(shape, axis=0) + if hasattr(self, "L") is False: + self.L = 1 + print( + "This is using self.L=1, which is likely not appropriate. Please calculate Jtor first to ensure the correct normalization." + ) return self.L * shape * self.Raxis def pressure(self, pn): @@ -1021,6 +1100,11 @@ def pressure(self, pn): ), axis=0, ) + if hasattr(self, "L") is False: + self.L = 1 + print( + "This is using self.L=1, which is likely not appropriate. Please calculate Jtor first to ensure the correct normalization." + ) pressure = self.L * norm_pressure * (self.psi_axis - self.psi_bndry) return pressure diff --git a/freegs4e/machine.py b/freegs4e/machine.py index c008ec5..c835b91 100644 --- a/freegs4e/machine.py +++ b/freegs4e/machine.py @@ -132,6 +132,33 @@ def createPsiGreens(self, R, Z): pgreen[label] = coil.createPsiGreens(R, Z) return pgreen + def createPsiGreensVec(self, R, Z): + """ + Calculate Greens functions + """ + pgreen = np.zeros(np.shape(R)) + for label, coil, multiplier in self.coils: + pgreen += multiplier * coil.createPsiGreens(R, Z) + return pgreen + + def createBrGreensVec(self, R, Z): + """ + Calculate Br Greens functions + """ + pgreen = np.zeros(np.shape(R)) + for label, coil, multiplier in self.coils: + pgreen += multiplier * coil.createBrGreensVec(R, Z) + return pgreen + + def createBzGreensVec(self, R, Z): + """ + Calculate Bz Greens functions + """ + pgreen = np.zeros(np.shape(R)) + for label, coil, multiplier in self.coils: + pgreen += multiplier * coil.createBzGreensVec(R, Z) + return pgreen + def calcPsiFromGreens(self, pgreen): """ Calculate poloidal flux from Greens functions. @@ -570,6 +597,24 @@ def createPsiGreens(self, R, Z): """ return self.controlPsi(R, Z) + def createPsiGreensVec(self, R, Z): + """ + Calculate Greens functions + """ + return self.controlPsi(R, Z) + + def createBrGreensVec(self, R, Z): + """ + Calculate Br Greens functions + """ + return self.controlBr(R, Z) + + def createBzGreensVec(self, R, Z): + """ + Calculate Bz Greens functions + """ + return self.controlBz(R, Z) + def calcPsiFromGreens(self, pgreen): """ Calculate poloidal flux from Greens functions. @@ -930,6 +975,16 @@ def __init__(self, coils, wall=None): self.coils = coils self.wall = wall + self.n_coils = len(coils) + self.current_vec = np.zeros(self.n_coils) + self.current_dummy_vec = np.zeros(self.n_coils) + self.getCurrentsVec() + + self.coil_names = list(self.getCurrents().keys()) + self.coil_order = {} + for i, coil in enumerate(self.coil_names): + self.coil_order[coil] = i + def __repr__(self): """ Return a string representation of the Machine object. @@ -1049,6 +1104,52 @@ def createPsiGreens(self, R, Z): pgreen[label] = coil.createPsiGreens(R, Z) return pgreen + def createPsiGreensVec(self, R, Z, coils=None): + """ + Pre-computes the Greens functions + and puts into arrays for each coil. This map can then be + called at a later time, and quickly return the field + """ + if coils == None: + coils = self.coils + + pgreen = np.zeros([len(coils)] + list(np.shape(R))) + i = 0 + for label, coil in coils: + pgreen[i] = coil.createPsiGreensVec(R, Z) + i += 1 + return pgreen + + def createBrGreensVec(self, R, Z, coils=None): + """ + Pre-computes the Br Greens functions + and puts into arrays for each coil. + """ + if coils == None: + coils = self.coils + + pgreen = np.zeros([len(coils)] + list(np.shape(R))) + i = 0 + for label, coil in coils: + pgreen[i] = coil.createBrGreensVec(R, Z) + i += 1 + return pgreen + + def createBzGreensVec(self, R, Z, coils=None): + """ + Pre-computes the Br Greens functions + and puts into arrays for each coil. + """ + if coils == None: + coils = self.coils + + pgreen = np.zeros([len(coils)] + list(np.shape(R))) + i = 0 + for label, coil in coils: + pgreen[i] = coil.createBzGreensVec(R, Z) + i += 1 + return pgreen + def calcPsiFromGreens(self, pgreen): """ Calculate poloidal flux from Greens functions. @@ -1069,6 +1170,15 @@ def calcPsiFromGreens(self, pgreen): psi_coils += coil.calcPsiFromGreens(pgreen[label]) return psi_coils + def getPsitokamak(self, vgreen): + """Uses self.current_vec and vgreen to calculate the plasma flux from all + coils, active and passive + """ + tokamak_psi = np.sum( + vgreen * self.current_vec[:, np.newaxis, np.newaxis], axis=0 + ) + return tokamak_psi + def Br(self, R, Z): """ Calculate radial magnetic field Br at (R,Z). @@ -1230,8 +1340,11 @@ def setControlCurrents(self, currents): """ controlcoils = [coil for label, coil in self.coils if coil.control] - for coil, current in zip(controlcoils, currents): - coil.current = current + controlcoils_labels = [ + label for label, coil in self.coils if coil.control + ] + for i in range(len(controlcoils_labels)): + self.set_coil_current(controlcoils_labels[i], currents[i]) def printCurrents(self): """ @@ -1284,6 +1397,49 @@ def getCurrents(self): currents[label] = coil.current return currents + def getCurrentsVec(self, coils=None): + """ + Returns an array of coil currents in Amps + """ + if coils == None: + coils = self.coils + + currents = np.zeros(len(coils)) + i = 0 + for label, coil in coils: + currents[i] = coil.current + i += 1 + self.current_vec = currents + return currents + + def set_coil_current(self, coil_label, current_value): + """Allows to set the current value of a coil + + Parameters + ---------- + coil_label : string + Coil label as from self.getCurrents().keys() + current_value : float + Current value in Amps + """ + i = self.coil_order[coil_label] + self.coils[i][1].current = current_value + self.current_vec[i] = current_value + + def set_all_coil_currents(self, current_vec): + """Allows to set all current values in the vec, + assuming correct order. + + Parameters + ---------- + current_vec : np.array + Current values in Amps + """ + + self.current_vec[: len(current_vec)] = current_vec + for i in range(len(current_vec)): + self.coils[i][1].current = current_vec[i] + def plot(self, axis=None, show=True): """ Plot the coils in the circuit. diff --git a/freegs4e/plotting.py b/freegs4e/plotting.py index 062f080..4c7cfb7 100644 --- a/freegs4e/plotting.py +++ b/freegs4e/plotting.py @@ -69,6 +69,53 @@ def plotConstraints(control, axis=None, show=True): return axis +def plotIOConstraints(control, axis=None, show=True): + """ + Plots constraints used for coil current control + + axis - Specify the axis on which to plot + show - Call matplotlib.pyplot.show() before returning + + """ + + import matplotlib + import matplotlib.pyplot as plt + + cmap = matplotlib.cm.get_cmap("gnuplot") + + if axis is None: + fig = plt.figure() + axis = fig.add_subplot(111) + + # Locations of the X-points + if control.null_points is not None: + axis.plot( + control.null_points[0], + control.null_points[1], + "1", + color="purple", + markersize=10, + ) + axis.plot( + [], [], "1", color="purple", markersize=10, label="Null-points" + ) + + # Isoflux surfaces + if control.isoflux_set is not None: + color = cmap(np.random.random()) + for i, isoflux in enumerate(control.isoflux_set): + axis.plot(isoflux[0], isoflux[1], "+", color=color, markersize=10) + axis.plot( + [], [], "+", color=color, label=f"Isoflux_{i}", markersize=10 + ) + + if show: + plt.legend(loc="upper right") + plt.show() + + return axis + + def plotEquilibrium( eq, axis=None, show=True, oxpoints=True, wall=True, limiter=True ): @@ -107,7 +154,7 @@ def plotEquilibrium( xpt = eq._profiles.xpt for r, z, _ in xpt: - axis.plot(r, z, "ro") + axis.plot(r, z, "rx") for r, z, _ in opt: axis.plot(r, z, "gx") @@ -142,10 +189,10 @@ def plotEquilibrium( ) # Add legend - axis.plot([], [], "ro", label="X-points") + axis.plot([], [], "rx", label="X-points") axis.plot([], [], "r", label="Separatrix") if opt is not []: - axis.plot([], [], "go", label="O-points") + axis.plot([], [], "gx", label="O-points") except: print(