diff --git a/doc/grid-file.rst b/doc/grid-file.rst index a0f38836..5b104039 100644 --- a/doc/grid-file.rst +++ b/doc/grid-file.rst @@ -132,6 +132,11 @@ BoutMesh writes three poloidal coordinates to the grid file: - The total poloidal distance around a closed flux surface in the core. Not calculated on open flux surfaces. + * - ``total_parallel_distance`` + + - The total distance parallel to a magnetic field line around a closed + flux surface in the core. Not calculated on open flux surfaces. + * - ``ShiftAngle`` - The total toroidal angular displacement when following a field line one @@ -259,6 +264,13 @@ Integral quantities surface to the grid point (on open field lines), or from the poloidal location of the lower X-point (on closed field lines). + * - ``parallel_distance`` + + - Distance (in metres) parallel to a magnetic field line from the lower + divertor target of each flux surface to the grid point (on open field + lines), or from the poloidal location of the lower X-point (on closed + field lines). + * - ``zShift`` - Toroidal displacement of a field line followed from some reference diff --git a/doc/whats-new.md b/doc/whats-new.md index 593b91a2..4c24b3cb 100644 --- a/doc/whats-new.md +++ b/doc/whats-new.md @@ -8,6 +8,7 @@ Release history - Set the dimension for R_closed_wall and Z_closed_wall to 'closed_wall'. Fixes loading of grid files by xBOUT (#191). By [John Omotani](https://github.com/johnomotani) +- Calculate the parallel distance along field lines, and save to the output (#193). ### New features diff --git a/hypnotoad/core/equilibrium.py b/hypnotoad/core/equilibrium.py index d9171da0..f9885f19 100644 --- a/hypnotoad/core/equilibrium.py +++ b/hypnotoad/core/equilibrium.py @@ -494,10 +494,11 @@ class FineContour: ), ) - def __init__(self, parentContour, settings, *, psi): + def __init__(self, parentContour, settings, *, psi, equilibrium): self.parentContour = parentContour self.user_options = self.user_options_factory.create(settings) self.distance = None + self.parallel_distance = None Nfine = self.user_options.finecontour_Nfine endInd = self.parentContour.endInd @@ -561,9 +562,9 @@ def __init__(self, parentContour, settings, *, psi): self.parentContour.endInd ].as_ndarray() - self.equaliseSpacing(psi=psi) + self.equaliseSpacing(psi=psi, equilibrium=equilibrium) - def extend(self, *, psi, extend_lower=0, extend_upper=0): + def extend(self, *, psi, equilibrium, extend_lower=0, extend_upper=0): Nfine = self.user_options.finecontour_Nfine parentCopy = self.parentContour.newContourFromSelf() @@ -646,9 +647,9 @@ def extend(self, *, psi, extend_lower=0, extend_upper=0): self.startInd = self.extend_lower_fine self.endInd = Nfine - 1 + self.extend_lower_fine - self.equaliseSpacing(psi=psi, reallocate=True) + self.equaliseSpacing(psi=psi, equilibrium=equilibrium, reallocate=True) - def equaliseSpacing(self, *, psi, reallocate=False): + def equaliseSpacing(self, *, psi, equilibrium, reallocate=False): """ Adjust the positions of points in this :class:`FineContour ` so they have a constant distance @@ -686,7 +687,7 @@ def equaliseSpacing(self, *, psi, reallocate=False): self.refine(psi=psi, skip_endpoints=True) - self.calcDistance(reallocate=reallocate) + self.calcDistance(reallocate=reallocate, equilibrium=equilibrium) ds = self.distance[1:] - self.distance[:-1] # want constant spacing, so ds has a constant value @@ -773,7 +774,7 @@ def equaliseSpacing(self, *, psi, reallocate=False): self.refine(psi=psi, skip_endpoints=True) - self.calcDistance() + self.calcDistance(equilibrium=equilibrium) ds = self.distance[1:] - self.distance[:-1] # want constant spacing, so ds has a constant value @@ -818,7 +819,7 @@ def equaliseSpacing(self, *, psi, reallocate=False): def totalDistance(self): return self.distance[self.endInd] - self.distance[self.startInd] - def calcDistance(self, *, reallocate=False): + def calcDistance(self, *, equilibrium, reallocate=False): """ Calculate poloidal distance from the start of this :class:`FineContour `. @@ -829,8 +830,35 @@ def calcDistance(self, *, reallocate=False): """ if self.distance is None or reallocate: self.distance = numpy.zeros(self.positions.shape[0]) + if self.parallel_distance is None or reallocate: + self.parallel_distance = numpy.zeros(self.positions.shape[0]) deltaSquared = (self.positions[1:] - self.positions[:-1]) ** 2 - self.distance[1:] = numpy.cumsum(numpy.sqrt(numpy.sum(deltaSquared, axis=1))) + delta_poloidal = numpy.sqrt(numpy.sum(deltaSquared, axis=1)) + self.distance[1:] = numpy.cumsum(delta_poloidal) + + if self.parentContour.psival is None: + # This is an EquilibriumRegion that is not necessarily following a single + # flux surface, so does not have a `psival`. We do not need the distance + # along this contour, so just set to NaN. + self.parallel_distance[:] = numpy.nan + else: + fine_contour_Bt = ( + equilibrium.fpol(self.parentContour.psival) / self.positions[:, 0] + ) + fine_contour_Br = equilibrium.Bp_R( + self.positions[:, 0], self.positions[:, 1] + ) + fine_contour_Bz = equilibrium.Bp_Z( + self.positions[:, 0], self.positions[:, 1] + ) + fine_contour_mod_Bp = numpy.sqrt(fine_contour_Br**2 + fine_contour_Bz**2) + fine_contour_B = numpy.sqrt(fine_contour_Bt**2 + fine_contour_mod_Bp**2) + midpoints_B = 0.5 * (fine_contour_B[:-1] + fine_contour_B[1:]) + midpoints_mod_Bp = 0.5 * ( + fine_contour_mod_Bp[:-1] + fine_contour_mod_Bp[1:] + ) + delta_parallel = midpoints_B / midpoints_mod_Bp * delta_poloidal + self.parallel_distance[1:] = numpy.cumsum(delta_parallel) def interpFunction(self, *, kind="linear"): distance = self.distance - self.distance[self.startInd] @@ -973,17 +1001,18 @@ def interpSSperp(self, vec, kind="linear"): return s_of_sperp, s_perp_total - def getDistance(self, p): + def getDistance(self, p, *, parallel=False): """ - Find the poloidal distance from the start of this contour of a point ``p``. + Find the poloidal and parallel distances from the start of this contour of a + point ``p``. Assume ``p`` is a point on the contour so has the correct psi-value. - Result is calculated as the weighted mean of the poloidal distances of the two - nearest points on the :class:`FineContour - ` (weighted by the relative distance - from ``p`` to each :class:`FineContour ` - point). + Result is calculated as the weighted mean of the poloidal (or parallel) + distances of the two nearest points on the :class:`FineContour + ` (weighted by the relative poloidal (in + both cases) distance from ``p`` to each :class:`FineContour + ` point). """ p = p.as_ndarray() @@ -1012,7 +1041,15 @@ def getDistance(self, p): # as their distances from the point r = d2 / (d1 + d2) - return r * self.distance[i1] + (1.0 - r) * self.distance[i2] + distance = r * self.distance[i1] + (1.0 - r) * self.distance[i2] + + # Weight by poloidal distances even when calculating parallel distance, + # because we cannot get the parallel displacement of `p`, only its position. + parallel_distance = ( + r * self.parallel_distance[i1] + (1.0 - r) * self.parallel_distance[i2] + ) + + return distance, parallel_distance def plot(self, *args, psi=None, ax=None, **kwargs): """ @@ -1084,6 +1121,7 @@ def __init__(self, *, points, psival, settings, Rrange, Zrange): self._fine_contour = None self._distance = None + self._parallel_distance = None # Value of vector potential on this contour self.psival = psival @@ -1106,6 +1144,7 @@ def _reset_cached(self): # Reset all cached objects/values because the contour has been changed self._fine_contour = None self._distance = None + self._parallel_distance = None @property def startInd(self): @@ -1153,7 +1192,7 @@ def extend_upper(self, val): self._reset_cached() self._extend_upper = val - def get_fine_contour(self, *, psi=None): + def get_fine_contour(self, *, psi=None, equilibrium=None): """ Get the FineContour associated with this PsiContour @@ -1163,15 +1202,21 @@ def get_fine_contour(self, *, psi=None): if self._fine_contour is None: if psi is None: raise ValueError("Poloidal flux psi needed to create FineContour") - self._fine_contour = FineContour(self, dict(self.user_options), psi=psi) + if equilibrium is None: + raise ValueError("equilibrium object needed to create FineContour") + self._fine_contour = FineContour( + self, dict(self.user_options), psi=psi, equilibrium=equilibrium + ) # Ensure that the fine contour is long enough - self.checkFineContourExtend(psi=psi) + self.checkFineContourExtend(psi=psi, equilibrium=equilibrium) return self._fine_contour - def get_distance(self, *, psi): + def get_distance(self, *, psi, equilibrium): if self._distance is None: - fine_contour = self.get_fine_contour(psi=psi) - self._distance = [fine_contour.getDistance(p) for p in self] + fine_contour = self.get_fine_contour(psi=psi, equilibrium=equilibrium) + both_distances = [fine_contour.getDistance(p) for p in self] + self._distance = [d[0] for d in both_distances] + self._parallel_distance = [d[1] for d in both_distances] d = numpy.array(self._distance) if not numpy.all(d[1:] - d[:-1] > 0.0): print("\nPsiContour distance", self._distance) @@ -1192,6 +1237,11 @@ def get_distance(self, *, psi): ) return self._distance + def get_parallel_distance(self, *, psi, equilibrium): + if self._parallel_distance is None: + self.get_distance(psi=psi, equilibrium=equilibrium) + return self._parallel_distance + def __iter__(self): return self.points.__iter__() @@ -1312,8 +1362,8 @@ def insertFindPosition(self, point): self.insert(minind, point) return minind - def totalDistance(self, *, psi): - distance = self.get_distance(psi=psi) + def totalDistance(self, *, psi, equilibrium): + distance = self.get_distance(psi=psi, equilibrium=equilibrium) return distance[self.endInd] - distance[self.startInd] def reverse(self): @@ -1600,8 +1650,8 @@ def getRefined(self, skip_endpoints=False, *, width=None, atol=None, **kwargs): return self.newContourFromSelf(points=newpoints) - def interpFunction(self, *, psi): - return self.get_fine_contour(psi=psi).interpFunction() + def interpFunction(self, *, psi, equilibrium): + return self.get_fine_contour(psi=psi, equilibrium=equilibrium).interpFunction() def _coarseInterp(self, *, kind="cubic"): distance = [0.0] @@ -1671,13 +1721,13 @@ def _coarseExtrapUpper(self, reference_ind, *, kind="cubic"): ) return lambda s: Point2D(interpR(s), interpZ(s)) - def contourSfunc(self, *, psi, kind="cubic"): + def contourSfunc(self, *, psi, equilibrium, kind="cubic"): """ Function interpolating distance as a function of index for the current state of this contour. When outside [startInd, endInd], set to constant so the results aren't affected by extrapolation errors. """ - distance = self.get_distance(psi=psi) + distance = self.get_distance(psi=psi, equilibrium=equilibrium) interpS = interpolate.interp1d( numpy.arange(len(self), dtype=float), distance, @@ -1703,7 +1753,7 @@ def contourSfunc(self, *, psi, kind="cubic"): ], ) - def interpSSperp(self, vec, *, psi): + def interpSSperp(self, vec, *, psi, equilibrium): """ Returns ------- @@ -1714,7 +1764,7 @@ def interpSSperp(self, vec, *, psi): contour. 2. the total perpendicular distance between startInd and endInd of the contour. """ - return self.get_fine_contour(psi=psi).interpSSperp(vec) + return self.get_fine_contour(psi=psi, equilibrium=equilibrium).interpSSperp(vec) def regrid(self, *args, **kwargs): """ @@ -1728,6 +1778,7 @@ def getRegridded( npoints, *, psi, + equilibrium, width=None, atol=None, sfunc=None, @@ -1777,7 +1828,7 @@ def getRegridded( # offset fine_contour.interpFunction in case sfunc(0.)!=0. sbegin = sfunc(0.0) else: - d = self.get_distance(psi=psi) + d = self.get_distance(psi=psi, equilibrium=equilibrium) s = (d[self.endInd] - d[self.startInd]) / (npoints - 1) * indices sbegin = 0.0 @@ -1787,7 +1838,7 @@ def getRegridded( # re-gridding should be the last change that is made to a PsiContour # Calling get_fine_contour() ensures that self._fine_contour has been initialised - self.get_fine_contour(psi=psi) + self.get_fine_contour(psi=psi, equilibrium=equilibrium) orig_extend_lower = self._fine_contour.extend_lower_fine orig_extend_upper = self._fine_contour.extend_upper_fine @@ -1798,7 +1849,9 @@ def getRegridded( while ( s[0] < -self._fine_contour.distance[self._fine_contour.startInd] - tol_lower ): - self._fine_contour.extend(extend_lower=max(orig_extend_lower, 1), psi=psi) + self._fine_contour.extend( + extend_lower=max(orig_extend_lower, 1), psi=psi, equilibrium=equilibrium + ) tol_upper = 0.25 * ( self._fine_contour.distance[-1] - self._fine_contour.distance[-2] @@ -1809,7 +1862,9 @@ def getRegridded( - self._fine_contour.distance[self._fine_contour.startInd] + tol_upper ): - self._fine_contour.extend(extend_upper=max(orig_extend_upper, 1), psi=psi) + self._fine_contour.extend( + extend_upper=max(orig_extend_upper, 1), psi=psi, equilibrium=equilibrium + ) interp_unadjusted = self._fine_contour.interpFunction() @@ -1838,13 +1893,13 @@ def interp(s): return new_contour - def checkFineContourExtend(self, *, psi): + def checkFineContourExtend(self, *, psi, equilibrium): """ Ensure that self._fine_contour extends past the first and last points of this PsiContour """ - fine_contour = self.get_fine_contour(psi=psi) + fine_contour = self.get_fine_contour(psi=psi, equilibrium=equilibrium) # check first point p = numpy.array([*self[0]]) @@ -1902,10 +1957,13 @@ def checkFineContourExtend(self, *, psi): return else: fine_contour.extend( - psi=psi, extend_lower=n_extend_lower, extend_upper=n_extend_upper + psi=psi, + equilibrium=equilibrium, + extend_lower=n_extend_lower, + extend_upper=n_extend_upper, ) # Call recursively to check extending has gone far enough - self.checkFineContourExtend(psi=psi) + self.checkFineContourExtend(psi=psi, equilibrium=equilibrium) def temporaryExtend( self, *, psi, extend_lower=0, extend_upper=0, ds_lower=None, ds_upper=None @@ -2689,7 +2747,7 @@ def nxInsideSeparatrix(self): def getRefined(self, *args, **kwargs): return self.newRegionFromPsiContour(super().getRefined(*args, **kwargs)) - def getRegridded(self, radialIndex, *, psi, **kwargs): + def getRegridded(self, radialIndex, *, psi, equilibrium, **kwargs): """ """ for wrong_argument in ["npoints", "extend_lower", "extend_upper", "sfunc"]: # these are valid arguments to PsiContour.getRegridded, but not to @@ -2708,7 +2766,7 @@ def getRegridded(self, radialIndex, *, psi, **kwargs): extend_upper = 2 * self.user_options.y_boundary_guards else: extend_upper = 0 - distance = self.get_distance(psi=psi) + distance = self.get_distance(psi=psi, equilibrium=equilibrium) sfunc = self.getSfuncFixedSpacing( 2 * self.ny_noguards + 1, distance[self.endInd] - distance[self.startInd], @@ -2717,6 +2775,7 @@ def getRegridded(self, radialIndex, *, psi, **kwargs): super().getRegridded( 2 * self.ny_noguards + 1, psi=psi, + equilibrium=equilibrium, extend_lower=extend_lower, extend_upper=extend_upper, sfunc=sfunc, @@ -2920,7 +2979,7 @@ def combineSfuncs( if vec_lower is None: sfunc_fixed_lower = self.getSfuncFixedSpacing( 2 * self.ny_noguards + 1, - contour.totalDistance(psi=self.psi), + contour.totalDistance(psi=self.psi, equilibrium=self.equilibrium), method="monotonic", spacing_lower=spacing_lower, spacing_upper=spacing_upper, @@ -2938,7 +2997,7 @@ def combineSfuncs( if vec_upper is None: sfunc_fixed_upper = self.getSfuncFixedSpacing( 2 * self.ny_noguards + 1, - contour.totalDistance(psi=self.psi), + contour.totalDistance(psi=self.psi, equilibrium=self.equilibrium), method="monotonic", spacing_lower=spacing_lower, spacing_upper=spacing_upper, @@ -3171,7 +3230,9 @@ def new_sfunc(i): (sfunc_fixed_upper, "fixed perp upper"), ], xind=contour.global_xind, - total_distance=contour.totalDistance(psi=self.psi), + total_distance=contour.totalDistance( + psi=self.psi, equilibrium=self.equilibrium + ), prefix="nonorthogonal_", ) except ValueError: @@ -3226,7 +3287,9 @@ def getSfuncFixedPerpSpacing( else: d_upper = spacing_upper - s_of_sperp, s_perp_total = contour.interpSSperp(surface_direction, psi=self.psi) + s_of_sperp, s_perp_total = contour.interpSSperp( + surface_direction, psi=self.psi, equilibrium=self.equilibrium + ) sperp_func = self.getMonotonicPoloidalDistanceFunc( s_perp_total, N - 1, N_norm, d_lower=d_lower, d_upper=d_upper ) diff --git a/hypnotoad/core/mesh.py b/hypnotoad/core/mesh.py index d303a350..b5221f60 100644 --- a/hypnotoad/core/mesh.py +++ b/hypnotoad/core/mesh.py @@ -289,7 +289,8 @@ def __init__( # Use self.equilibriumRegion.fine_contour for the vector along the separatrix # because then the vector will not change when the grid resolution changes fine_contour = self.equilibriumRegion.get_fine_contour( - psi=self.equilibriumRegion.psi + psi=self.equilibriumRegion.psi, + equilibrium=self.meshParent.equilibrium, ) if self.equilibriumRegion.wallSurfaceAtStart is None: # lower end @@ -385,7 +386,9 @@ def addPointAtWallToContours(self): # point makes the spacing of points on the contour not-smooth) and adjusted for # the change in distance after redefining startInd to be at the wall self.sfunc_orthogonal_list = [ - contour.contourSfunc(psi=self.equilibriumRegion.psi) + contour.contourSfunc( + psi=self.equilibriumRegion.psi, equilibrium=self.meshParent.equilibrium + ) for contour in self.contours ] @@ -447,7 +450,8 @@ def correct_sfunc_orthogonal_and_set_startInd( # Calculate total distance using a new FineContour, created at this # point new_total_distance = contour.totalDistance( - psi=self.equilibriumRegion.psi + psi=self.equilibriumRegion.psi, + equilibrium=self.meshParent.equilibrium, ) return ( @@ -647,6 +651,7 @@ def get_sfunc(i_contour, contour, sfunc_orthogonal): c.regrid( 2 * self.ny_noguards + 1, psi=self.equilibriumRegion.psi, + equilibrium=self.meshParent.equilibrium, sfunc=get_sfunc(i, c, sfunc_orth), extend_lower=self.equilibriumRegion.extend_lower, extend_upper=self.equilibriumRegion.extend_upper, @@ -856,6 +861,7 @@ def geometry1(self): self.Bpxy = numpy.sqrt(self.Brxy**2 + self.Bzxy**2) self.calcPoloidalDistance() + self.calcParallelDistance() if hasattr( self.meshParent.equilibrium.regions[self.equilibriumRegion.name], "pressure" @@ -1363,13 +1369,19 @@ def calcHy(self): for i in range(self.nx): print(f"{self.name} calcHy {i} / {2 * self.nx + 1}", end="\r", flush=True) d = numpy.array( - self.contours[2 * i + 1].get_distance(psi=self.equilibriumRegion.psi) + self.contours[2 * i + 1].get_distance( + psi=self.equilibriumRegion.psi, + equilibrium=self.meshParent.equilibrium, + ) ) hy.centre[i, :] = d[2::2] - d[:-2:2] hy.ylow[i, 1:-1] = d[3:-1:2] - d[1:-3:2] if self.connections["lower"] is not None: cbelow = self.getNeighbour("lower").contours[2 * i + 1] - dbelow = cbelow.get_distance(psi=self.equilibriumRegion.psi) + dbelow = cbelow.get_distance( + psi=self.equilibriumRegion.psi, + equilibrium=self.meshParent.equilibrium, + ) hy.ylow[i, 0] = d[1] - d[0] + dbelow[-1] - dbelow[-2] else: # no region below, so estimate distance to point before '0' as the same @@ -1377,7 +1389,10 @@ def calcHy(self): hy.ylow[i, 0] = 2.0 * (d[1] - d[0]) if self.connections["upper"] is not None: cabove = self.getNeighbour("upper").contours[2 * i + 1] - dabove = cabove.get_distance(psi=self.equilibriumRegion.psi) + dabove = cabove.get_distance( + psi=self.equilibriumRegion.psi, + equilibrium=self.meshParent.equilibrium, + ) hy.ylow[i, -1] = d[-1] - d[-2] + dabove[1] - dabove[0] else: # no region below, so estimate distance to point before '0' as the same @@ -1391,13 +1406,19 @@ def calcHy(self): flush=True, ) d = numpy.array( - self.contours[2 * i].get_distance(psi=self.equilibriumRegion.psi) + self.contours[2 * i].get_distance( + psi=self.equilibriumRegion.psi, + equilibrium=self.meshParent.equilibrium, + ) ) hy.xlow[i, :] = d[2::2] - d[:-2:2] hy.corners[i, 1:-1] = d[3:-1:2] - d[1:-3:2] if self.connections["lower"] is not None: cbelow = self.getNeighbour("lower").contours[2 * i] - dbelow = cbelow.get_distance(psi=self.equilibriumRegion.psi) + dbelow = cbelow.get_distance( + psi=self.equilibriumRegion.psi, + equilibrium=self.meshParent.equilibrium, + ) hy.corners[i, 0] = d[1] - d[0] + dbelow[-1] - dbelow[-2] else: # no region below, so estimate distance to point before '0' as the same @@ -1405,7 +1426,10 @@ def calcHy(self): hy.corners[i, 0] = 2.0 * (d[1] - d[0]) if self.connections["upper"] is not None: cabove = self.getNeighbour("upper").contours[2 * i] - dabove = cabove.get_distance(psi=self.equilibriumRegion.psi) + dabove = cabove.get_distance( + psi=self.equilibriumRegion.psi, + equilibrium=self.meshParent.equilibrium, + ) hy.corners[i, -1] = d[-1] - d[-2] + dabove[1] - dabove[0] else: # no region below, so estimate distance to point before '0' as the same @@ -1548,7 +1572,10 @@ def calcZShift(self): while True: print("calcZShift", region.name, end="\r", flush=True) for i, contour in enumerate(region.contours): - fine_contour = contour.get_fine_contour(psi=self.equilibriumRegion.psi) + fine_contour = contour.get_fine_contour( + psi=self.equilibriumRegion.psi, + equilibrium=self.meshParent.equilibrium, + ) fine_distance = fine_contour.distance def integrand_func(R, Z): @@ -1581,7 +1608,10 @@ def integrand_func(R, Z): fine_distance, zShift_fine, kind="linear", assume_sorted=True ) zShift_contour = zShift_interpolator( - contour.get_distance(psi=self.equilibriumRegion.psi) + contour.get_distance( + psi=self.equilibriumRegion.psi, + equilibrium=self.meshParent.equilibrium, + ) ) if i % 2 == 0: @@ -1650,21 +1680,25 @@ def calcPoloidalDistance(self): c = region.contours[2 * i + 1] # Cell-centre points region.poloidal_distance.centre[i, :] -= c.get_distance( - psi=self.meshParent.equilibrium.psi + psi=self.meshParent.equilibrium.psi, + equilibrium=self.meshParent.equilibrium, )[c.startInd] # ylow points region.poloidal_distance.ylow[i, :] -= c.get_distance( - psi=self.meshParent.equilibrium.psi + psi=self.meshParent.equilibrium.psi, + equilibrium=self.meshParent.equilibrium, )[c.startInd] for i in range(self.nx + 1): c = region.contours[2 * i] # Cell-centre points region.poloidal_distance.xlow[i, :] -= c.get_distance( - psi=self.meshParent.equilibrium.psi + psi=self.meshParent.equilibrium.psi, + equilibrium=self.meshParent.equilibrium, )[c.startInd] # ylow points region.poloidal_distance.corners[i, :] -= c.get_distance( - psi=self.meshParent.equilibrium.psi + psi=self.meshParent.equilibrium.psi, + equilibrium=self.meshParent.equilibrium, )[c.startInd] # Get distances from contours @@ -1673,21 +1707,25 @@ def calcPoloidalDistance(self): c = region.contours[2 * i + 1] # Cell-centre points region.poloidal_distance.centre[i, :] += c.get_distance( - psi=self.meshParent.equilibrium.psi + psi=self.meshParent.equilibrium.psi, + equilibrium=self.meshParent.equilibrium, )[1::2] # ylow points region.poloidal_distance.ylow[i, :] += c.get_distance( - psi=self.meshParent.equilibrium.psi + psi=self.meshParent.equilibrium.psi, + equilibrium=self.meshParent.equilibrium, )[::2] for i in range(self.nx + 1): c = region.contours[2 * i] # Cell-centre points region.poloidal_distance.xlow[i, :] += c.get_distance( - psi=self.meshParent.equilibrium.psi + psi=self.meshParent.equilibrium.psi, + equilibrium=self.meshParent.equilibrium, )[1::2] # ylow points region.poloidal_distance.corners[i, :] += c.get_distance( - psi=self.meshParent.equilibrium.psi + psi=self.meshParent.equilibrium.psi, + equilibrium=self.meshParent.equilibrium, )[::2] next_region = region.getNeighbour("upper") @@ -1726,6 +1764,115 @@ def calcPoloidalDistance(self): :, -1 ] + def calcParallelDistance(self): + """ + Calculate parallel distance by following contours between regions. + """ + # Cannot just test 'connections['lower'] is not None' because periodic regions + # always have a lower connection - requires us to give a yGroupIndex to each + # region when creating the groups. + if self.yGroupIndex != 0: + return None + + region = self + region.parallel_distance = MultiLocationArray(region.nx, region.ny) + region.parallel_distance.centre = 0.0 + region.parallel_distance.ylow = 0.0 + region.parallel_distance.xlow = 0.0 + region.parallel_distance.corners = 0.0 + + # Initialise so that distance counts from the lower wall (for SOL/PFR) or wall + # (for core) + for i in range(self.nx): + c = region.contours[2 * i + 1] + # Cell-centre points + region.parallel_distance.centre[i, :] -= c.get_parallel_distance( + psi=self.meshParent.equilibrium.psi, + equilibrium=self.meshParent.equilibrium, + )[c.startInd] + # ylow points + region.parallel_distance.ylow[i, :] -= c.get_parallel_distance( + psi=self.meshParent.equilibrium.psi, + equilibrium=self.meshParent.equilibrium, + )[c.startInd] + for i in range(self.nx + 1): + c = region.contours[2 * i] + # Cell-centre points + region.parallel_distance.xlow[i, :] -= c.get_parallel_distance( + psi=self.meshParent.equilibrium.psi, + equilibrium=self.meshParent.equilibrium, + )[c.startInd] + # ylow points + region.parallel_distance.corners[i, :] -= ( + c.get_parallel_distance( + psi=self.meshParent.equilibrium.psi, + equilibrium=self.meshParent.equilibrium, + )[c.startInd], + ) + + # Get distances from contours + while True: + for i in range(self.nx): + c = region.contours[2 * i + 1] + # Cell-centre points + region.parallel_distance.centre[i, :] += c.get_parallel_distance( + psi=self.meshParent.equilibrium.psi, + equilibrium=self.meshParent.equilibrium, + )[1::2] + # ylow points + region.parallel_distance.ylow[i, :] += c.get_parallel_distance( + psi=self.meshParent.equilibrium.psi, + equilibrium=self.meshParent.equilibrium, + )[::2] + for i in range(self.nx + 1): + c = region.contours[2 * i] + # Cell-centre points + region.parallel_distance.xlow[i, :] += c.get_parallel_distance( + psi=self.meshParent.equilibrium.psi, + equilibrium=self.meshParent.equilibrium, + )[1::2] + # ylow points + region.parallel_distance.corners[i, :] += c.get_parallel_distance( + psi=self.meshParent.equilibrium.psi, + equilibrium=self.meshParent.equilibrium, + )[::2] + + next_region = region.getNeighbour("upper") + if (next_region is None) or (next_region is self): + # Note: If periodic, next_region is self (back to start) + break + else: + # Initialise with values at the lower y-boundary of next_region + next_region.parallel_distance = MultiLocationArray( + next_region.nx, next_region.ny + ) + next_region.parallel_distance.centre[:, :] = ( + region.parallel_distance.ylow[:, -1, numpy.newaxis] + ) + next_region.parallel_distance.ylow[:, :] = ( + region.parallel_distance.ylow[:, -1, numpy.newaxis] + ) + next_region.parallel_distance.xlow[:, :] = ( + region.parallel_distance.corners[:, -1, numpy.newaxis] + ) + next_region.parallel_distance.corners[:, :] = ( + region.parallel_distance.corners[:, -1, numpy.newaxis] + ) + region = next_region + + # Save total parallel distance in core + self.total_parallel_distance = MultiLocationArray(region.nx, 1) + if self.connections["lower"] is not None: + # This is a periodic region (we already checked that the self.yGroupIndex is + # 0). + # 'region' is the last region in the y-group + self.total_parallel_distance.centre[:, 0] = region.parallel_distance.ylow[ + :, -1 + ] + self.total_parallel_distance.xlow[:, 0] = region.parallel_distance.corners[ + :, -1 + ] + def getNeighbour(self, face): if self.connections[face] is None: return None @@ -2306,13 +2453,13 @@ def smooth_mask(mark): setattr(self, varname, tmp) -def _calc_contour_distance(i, c, *, psi, **kwargs): +def _calc_contour_distance(i, c, *, psi, equilibrium, **kwargs): print( f"Calculating contour distances: {i + 1}", end="\r", flush=True, ) - c.get_distance(psi=psi) + c.get_distance(psi=psi, equilibrium=equilibrium) return c @@ -2369,7 +2516,7 @@ def _find_intersection( break count = 0 - distance = contour.get_distance(psi=psi) + distance = contour.get_distance(psi=psi, equilibrium=equilibrium) ds_extend = distance[1] - distance[0] while coarse_lower_intersect is None: # contour has not yet intersected with wall, so make it longer and @@ -2401,8 +2548,8 @@ def _find_intersection( # points # # first find nearest FineContour points - fine_contour = contour.get_fine_contour(psi=psi) - d = fine_contour.getDistance(coarse_lower_intersect) + fine_contour = contour.get_fine_contour(psi=psi, equilibrium=equilibrium) + d, _ = fine_contour.getDistance(coarse_lower_intersect) i_fine = numpy.searchsorted(fine_contour.distance, d) # Intersection should be between i_fine-1 and i_fine, but check # intervals on either side if necessary @@ -2455,7 +2602,7 @@ def _find_intersection( break count = 0 - distance = contour.get_distance(psi=psi) + distance = contour.get_distance(psi=psi, equilibrium=equilibrium) ds_extend = distance[-1] - distance[-2] while coarse_upper_intersect is None: # contour has not yet intersected with wall, so make it longer and @@ -2487,8 +2634,8 @@ def _find_intersection( # points # # first find nearest FineContour points - fine_contour = contour.get_fine_contour(psi=psi) - d = fine_contour.getDistance(coarse_upper_intersect) + fine_contour = contour.get_fine_contour(psi=psi, equilibrium=equilibrium) + d, _ = fine_contour.getDistance(coarse_upper_intersect) i_fine = numpy.searchsorted(fine_contour.distance, d) # Intersection should be between i_fine-1 and i_fine, but check # intervals on either side if necessary @@ -2525,7 +2672,7 @@ def _find_intersection( ) # Create FineContour for result, so that this is done in parallel - contour.get_fine_contour(psi=psi) + contour.get_fine_contour(psi=psi, equilibrium=equilibrium) return ( contour, @@ -2536,7 +2683,7 @@ def _find_intersection( ) -def _refine_extend(i, contour, *, psi, **kwargs): +def _refine_extend(i, contour, *, psi, equilibrium, **kwargs): print( "refine and extend", i, @@ -2544,7 +2691,7 @@ def _refine_extend(i, contour, *, psi, **kwargs): flush=True, ) contour.refine(psi=psi) - contour.checkFineContourExtend(psi=psi) + contour.checkFineContourExtend(psi=psi, equilibrium=equilibrium) return contour @@ -2695,6 +2842,7 @@ def makeRegions(self, parallel_map): eq_region_with_boundaries = eq_region.getRegridded( radialIndex=i, psi=self.equilibrium.psi, + equilibrium=self.equilibrium, width=self.user_options.refine_width, ) self.regions[region_id] = MeshRegion( @@ -3614,6 +3762,8 @@ def addFromRegionsXArray(name): addFromRegions("dy") addFromRegions("poloidal_distance") addFromRegionsXArray("total_poloidal_distance") + addFromRegions("parallel_distance") + addFromRegionsXArray("total_parallel_distance") addFromRegions("Brxy") addFromRegions("Bzxy") addFromRegions("Bpxy") diff --git a/hypnotoad/test_suite/test_equilibrium.py b/hypnotoad/test_suite/test_equilibrium.py index ba8e1bb1..d8de79bb 100644 --- a/hypnotoad/test_suite/test_equilibrium.py +++ b/hypnotoad/test_suite/test_equilibrium.py @@ -362,6 +362,24 @@ def test_closest_approach_right(): assert numpy.isclose(cpa, 0.5) +class ThisEquilibrium(Equilibrium): + def __init__(self, settings=None, wall=None): + if settings is None: + settings = {} + self.user_options = Equilibrium.user_options_factory.add( + refine_width=1.0e-5, refine_atol=2.0e-8 + ).create(settings) + + if wall is not None: + self.wall = wall + + self.fpol = lambda psi: 1.0 + 0.0 * psi + self.Bp_R = lambda R, Z: 1.0 + 0.0 * R + self.Bp_Z = lambda R, Z: 1.0 + 0.0 * R + + super().__init__({}) + + class TestContour: @pytest.fixture def testcontour(self): @@ -394,9 +412,22 @@ def psifunc(R, Z): return returnObject() - def test_distance(self, testcontour): + @pytest.fixture + def eq(self): + wall = [ + Point2D(-1.0, -1.0), + Point2D(1.0, -1.0), + Point2D(1.0, 1.0), + Point2D(-1.0, 1.0), + ] + eq = ThisEquilibrium(wall=wall) + return eq + + def test_distance(self, testcontour, eq): segment_length = testcontour.r * numpy.pi / (testcontour.npoints - 1) - assert testcontour.c.get_distance(psi=testcontour.psi) == pytest.approx( + assert testcontour.c.get_distance( + psi=testcontour.psi, equilibrium=eq + ) == pytest.approx( segment_length * numpy.arange(testcontour.npoints), abs=1.0e-4 ) @@ -412,29 +443,29 @@ def test_getitem(self, testcontour): assert p.R == tight_approx(testcontour.R[5]) assert p.Z == tight_approx(testcontour.Z[5]) - def test_append(self, testcontour): + def test_append(self, testcontour, eq): c = testcontour.c point_to_add = c[-2] - expected_distance = c.get_distance(psi=testcontour.psi)[-2] + expected_distance = c.get_distance(psi=testcontour.psi, equilibrium=eq)[-2] del c.points[-1] del c.points[-1] c.endInd = -1 c.append(point_to_add) - assert c.get_distance(psi=testcontour.psi)[-1] == pytest.approx( + assert c.get_distance(psi=testcontour.psi, equilibrium=eq)[-1] == pytest.approx( expected_distance, abs=4.0e-5 ) - def test_prepend(self, testcontour): + def test_prepend(self, testcontour, eq): c = testcontour.c c.startInd = 3 c.endInd = -1 point_to_add = c[1] - expected_distance = c.get_distance(psi=testcontour.psi)[1] + expected_distance = c.get_distance(psi=testcontour.psi, equilibrium=eq)[1] del c.points[0] del c.points[0] c.startInd = 1 c.prepend(point_to_add) - assert c.get_distance(psi=testcontour.psi)[1] == pytest.approx( + assert c.get_distance(psi=testcontour.psi, equilibrium=eq)[1] == pytest.approx( expected_distance, abs=6.0e-4 ) assert c.startInd == 1 @@ -455,20 +486,20 @@ def test_replace(self, testcontour): assert c[5] == point_to_add assert c.endInd == 22 - def test_reverse(self, testcontour): + def test_reverse(self, testcontour, eq): c = testcontour.c orig = deepcopy(c) c.reverse() n = len(orig) - total_d = orig.get_distance(psi=testcontour.psi)[-1] + total_d = orig.get_distance(psi=testcontour.psi, equilibrium=eq)[-1] for i in range(n): assert orig[n - 1 - i].R == tight_approx(c[i].R) assert orig[n - 1 - i].Z == tight_approx(c[i].Z) - assert total_d - orig.get_distance(psi=testcontour.psi)[ + assert total_d - orig.get_distance(psi=testcontour.psi, equilibrium=eq)[ n - 1 - i - ] == tight_approx(c.get_distance(psi=testcontour.psi)[i]) + ] == tight_approx(c.get_distance(psi=testcontour.psi, equilibrium=eq)[i]) def test_insertFindPosition(self, testcontour): c = testcontour.c @@ -500,8 +531,8 @@ def test_coarseInterp(self, testcontour): assert pend.R == pytest.approx(testcontour.R0 - testcontour.r, abs=1.0e-9) assert pend.Z == pytest.approx(testcontour.Z0, abs=1.0e-5) - def test_interpFunction(self, testcontour): - f = testcontour.c.interpFunction(psi=testcontour.psi) + def test_interpFunction(self, testcontour, eq): + f = testcontour.c.interpFunction(psi=testcontour.psi, equilibrium=eq) pstart = f(0.0) pend = f(numpy.pi * testcontour.r) assert pstart.R == pytest.approx(testcontour.R0 + testcontour.r, abs=1.0e-9) @@ -509,7 +540,7 @@ def test_interpFunction(self, testcontour): assert pend.R == pytest.approx(testcontour.R0 - testcontour.r, abs=1.0e-8) assert pend.Z == pytest.approx(testcontour.Z0, abs=1.0e-5) - def test_getRegridded(self, testcontour): + def test_getRegridded(self, testcontour, eq): orig = testcontour.c r = testcontour.r @@ -521,7 +552,7 @@ def sfunc_true(i): def sfunc(i): return ( numpy.sqrt(i / (newNpoints - 1)) - * orig.get_distance(psi=testcontour.psi)[-1] + * orig.get_distance(psi=testcontour.psi, equilibrium=eq)[-1] ) newTheta = sfunc_true(numpy.arange(newNpoints)) / r @@ -529,19 +560,20 @@ def sfunc(i): newZ = testcontour.Z0 + r * numpy.sin(newTheta) new = orig.getRegridded( - newNpoints, psi=testcontour.psi, sfunc=sfunc, width=1.0e-3 + newNpoints, psi=testcontour.psi, equilibrium=eq, sfunc=sfunc, width=1.0e-3 ) assert [p.R for p in new] == pytest.approx(newR, abs=4.0e-4) assert [p.Z for p in new] == pytest.approx(newZ, abs=4.0e-4) - def test_getRegridded_extend(self, testcontour): + def test_getRegridded_extend(self, testcontour, eq): c = testcontour.c orig = c.newContourFromSelf() new = c.getRegridded( testcontour.npoints, psi=testcontour.psi, + equilibrium=eq, width=0.1, extend_lower=1, extend_upper=2, @@ -564,7 +596,7 @@ def test_getRegridded_extend(self, testcontour): [orig[-3].R, 2.0 * testcontour.Z0 - orig[-3].Z], abs=1.0e-7 ) - def test_contourSfunc(self, testcontour): + def test_contourSfunc(self, testcontour, eq): c = testcontour.c c.startInd = 2 c.endInd = len(c) - 2 @@ -572,22 +604,24 @@ def test_contourSfunc(self, testcontour): c.extend_lower = 2 c.extend_upper = 2 - f = c.contourSfunc(psi=testcontour.psi) + f = c.contourSfunc(psi=testcontour.psi, equilibrium=eq) indices = numpy.arange(n, dtype=float) assert f(indices) == tight_approx( [ - d - c.get_distance(psi=testcontour.psi)[c.startInd] - for d in c.get_distance(psi=testcontour.psi)[c.startInd : c.endInd + 1] + d - c.get_distance(psi=testcontour.psi, equilibrium=eq)[c.startInd] + for d in c.get_distance(psi=testcontour.psi, equilibrium=eq)[ + c.startInd : c.endInd + 1 + ] ] ) assert f(-1.0) == tight_approx(0.0) assert f(n + 1.0) == tight_approx( - c.get_distance(psi=testcontour.psi)[c.endInd] - - c.get_distance(psi=testcontour.psi)[c.startInd] + c.get_distance(psi=testcontour.psi, equilibrium=eq)[c.endInd] + - c.get_distance(psi=testcontour.psi, equilibrium=eq)[c.startInd] ) - def test_contourSfunc_list(self, testcontour): + def test_contourSfunc_list(self, testcontour, eq): c1 = testcontour.c c1.startInd = 2 c1.endInd = len(c1) - 2 @@ -626,7 +660,7 @@ def test_contourSfunc_list(self, testcontour): # the last value from the loop instead of the value when 'sfunc' was added to # 'sfunc_list' for c in c_list: - sfunc_orig = c.contourSfunc(psi=testcontour.psi) + sfunc_orig = c.contourSfunc(psi=testcontour.psi, equilibrium=eq) sfunc_list.append(lambda i: sfunc_orig(i) - 3.0) @@ -643,7 +677,7 @@ def test_contourSfunc_list(self, testcontour): # This version does work, because when the lambda is evaluated it uses # 'sfunc_orig' from the scope of 'shift_sfunc' in which it was created. def shift_sfunc(c): - sfunc_orig = c.contourSfunc(psi=testcontour.psi) + sfunc_orig = c.contourSfunc(psi=testcontour.psi, equilibrium=eq) return lambda i: sfunc_orig(i) - 3.0 for c in c_list: @@ -656,24 +690,29 @@ def shift_sfunc(c): n2 / (npoints2 - 1.0) * 1.5 * numpy.pi * r - 3.0, abs=4.0e-6 ) - def test_interpSSperp(self, testcontour): + def test_interpSSperp(self, testcontour, eq): c = testcontour.c # Make c.startInd > 0 c.insert(0, Point2D(c[1].R, 2.0 * c[0].Z - c[1].Z)) # 'vec' argument is in Z-direction, so 's_perp' is displacement in R-direction - sfunc, s_perp_total = c.interpSSperp([0.0, 1.0], psi=testcontour.psi) + sfunc, s_perp_total = c.interpSSperp( + [0.0, 1.0], psi=testcontour.psi, equilibrium=eq + ) assert sfunc(0.0) == tight_approx(0.0) assert sfunc(1.0) == pytest.approx(numpy.pi / 2.0, abs=1.0e-6) assert sfunc(2.0) == pytest.approx(numpy.pi, abs=2.0e-6) assert s_perp_total == tight_approx(2.0) - def test_FineContour(self, testcontour): + def test_FineContour(self, testcontour, eq): testcontour.c.refine_width = 1.0e-2 fc = FineContour( - testcontour.c, dict(testcontour.c.user_options), psi=testcontour.psi + testcontour.c, + dict(testcontour.c.user_options), + psi=testcontour.psi, + equilibrium=eq, ) assert fc.totalDistance() == pytest.approx(numpy.pi, abs=1.0e-5) @@ -689,7 +728,7 @@ def test_FineContour(self, testcontour): testcontour.Z0 + r * numpy.sin(theta), abs=1.0e-4 ) - def test_finecontour_extent_lower(self, testcontour): + def test_finecontour_extent_lower(self, testcontour, eq): contour = testcontour.c contour_initial = contour[0] @@ -698,15 +737,15 @@ def test_finecontour_extent_lower(self, testcontour): # This will create contour._fine_contour and check that contour._distance # is monotonic at this point - contour.get_distance(psi=testcontour.psi) + contour.get_distance(psi=testcontour.psi, equilibrium=eq) contour.prepend(contour_initial) # Check that after modifying contour, contour._fine_contour still extends far # enough to give a monotonic contour._distance - contour.get_distance(psi=testcontour.psi) + contour.get_distance(psi=testcontour.psi, equilibrium=eq) - def test_finecontour_extent_upper(self, testcontour): + def test_finecontour_extent_upper(self, testcontour, eq): contour = testcontour.c contour_final = contour[-1] @@ -715,26 +754,12 @@ def test_finecontour_extent_upper(self, testcontour): # This will create contour._fine_contour and check that contour._distance # is monotonic at this point - contour.get_distance(psi=testcontour.psi) + contour.get_distance(psi=testcontour.psi, equilibrium=eq) contour.append(contour_final) # Check that after modifying contour, contour._fine_contour still extends far # enough to give a monotonic contour._distance - contour.get_distance(psi=testcontour.psi) - - -class ThisEquilibrium(Equilibrium): - def __init__(self, settings=None, wall=None): - if settings is None: - settings = {} - self.user_options = Equilibrium.user_options_factory.add( - refine_width=1.0e-5, refine_atol=2.0e-8 - ).create(settings) - - if wall is not None: - self.wall = wall - - super().__init__({}) + contour.get_distance(psi=testcontour.psi, equilibrium=eq) class TestEquilibrium: @@ -1281,7 +1306,7 @@ def test_getSqrtPoloidalDistanceFuncBothBoth(self, eqReg): def test_combineSfuncsPoloidalSpacing(self, eqReg): n = len(eqReg) - L = eqReg.totalDistance(psi=eqReg.psi) + L = eqReg.totalDistance(psi=eqReg.psi, equilibrium=eqReg.equilibrium) eqReg.resetNonorthogonalOptions( {"nonorthogonal_target_all_poloidal_spacing_length": L} @@ -1300,7 +1325,7 @@ def sfunc_orthogonal(i): def test_combineSfuncsPoloidalSpacingRangeLower(self, eqReg): n = len(eqReg) - L = eqReg.totalDistance(psi=eqReg.psi) + L = eqReg.totalDistance(psi=eqReg.psi, equilibrium=eqReg.equilibrium) new_settings = { "nonorthogonal_target_all_poloidal_spacing_length": L * 40.0 / (n - 1), @@ -1327,7 +1352,7 @@ def sfunc_orthogonal(i): def test_combineSfuncsPoloidalSpacingRangeUpper(self, eqReg): n = len(eqReg) - L = eqReg.totalDistance(psi=eqReg.psi) + L = eqReg.totalDistance(psi=eqReg.psi, equilibrium=eqReg.equilibrium) new_settings = { "nonorthogonal_target_all_poloidal_spacing_length": L * 40.0 / (n - 1), @@ -1354,7 +1379,7 @@ def sfunc_orthogonal(i): def test_combineSfuncsPoloidalSpacingRangeBoth(self, eqReg): n = len(eqReg) - L = eqReg.totalDistance(psi=eqReg.psi) + L = eqReg.totalDistance(psi=eqReg.psi, equilibrium=eqReg.equilibrium) new_settings = { "nonorthogonal_target_all_poloidal_spacing_length": L * 40.0 / (n - 1), @@ -1397,12 +1422,16 @@ def test_combineSfuncsPerpSpacingIntegrated(self, eqReg): eqReg.sin_angle_at_end = 1.0 eqReg.name = "inner_lower" - sfunc_orthogonal_original = eqReg.contourSfunc(psi=eqReg.psi) + sfunc_orthogonal_original = eqReg.contourSfunc( + psi=eqReg.psi, equilibrium=eqReg.equilibrium + ) # as if lower_wall intersect_index = 2 original_start = eqReg.startInd - distance_at_original_start = eqReg.get_distance(psi=eqReg.psi)[original_start] + distance_at_original_start = eqReg.get_distance( + psi=eqReg.psi, equilibrium=eqReg.equilibrium + )[original_start] d = 1.5 * 3.0 / (n - 1.0) eqReg.points.insert(intersect_index, Point2D(d, d)) @@ -1410,7 +1439,9 @@ def test_combineSfuncsPerpSpacingIntegrated(self, eqReg): if original_start >= intersect_index: original_start += 1 - distance_at_wall = eqReg.get_distance(psi=eqReg.psi)[intersect_index] + distance_at_wall = eqReg.get_distance( + psi=eqReg.psi, equilibrium=eqReg.equilibrium + )[intersect_index] sfunc_orthogonal = ( lambda i: sfunc_orthogonal_original(i) @@ -1422,12 +1453,14 @@ def test_combineSfuncsPerpSpacingIntegrated(self, eqReg): eqReg.extend_lower = intersect_index eqReg.extend_upper = 1 - d = eqReg.totalDistance(psi=eqReg.psi) + d = eqReg.totalDistance(psi=eqReg.psi, equilibrium=eqReg.equilibrium) sfunc = eqReg.combineSfuncs(eqReg, sfunc_orthogonal, [0.0, 1.0], [1.0, 0.0]) assert sfunc(0.0) == tight_approx(0.0) - assert sfunc(n - 1.0) == tight_approx(eqReg.totalDistance(psi=eqReg.psi)) + assert sfunc(n - 1.0) == tight_approx( + eqReg.totalDistance(psi=eqReg.psi, equilibrium=eqReg.equilibrium) + ) def test_combineSfuncsPoloidalSpacingIntegrated(self, eqReg): # This test follows roughly the operations in @@ -1443,7 +1476,9 @@ def test_combineSfuncsPoloidalSpacingIntegrated(self, eqReg): eqReg.ny_total = 40 eqReg.name = "outer_lower" - sfunc_orthogonal_original = eqReg.contourSfunc(psi=eqReg.psi) + sfunc_orthogonal_original = eqReg.contourSfunc( + psi=eqReg.psi, equilibrium=eqReg.equilibrium + ) # as if lower_wall intersect_index = 2 @@ -1454,9 +1489,13 @@ def test_combineSfuncsPoloidalSpacingIntegrated(self, eqReg): if original_start >= intersect_index: original_start += 1 - distance_at_original_start = eqReg.get_distance(psi=eqReg.psi)[original_start] + distance_at_original_start = eqReg.get_distance( + psi=eqReg.psi, equilibrium=eqReg.equilibrium + )[original_start] - distance_at_wall = eqReg.get_distance(psi=eqReg.psi)[intersect_index] + distance_at_wall = eqReg.get_distance( + psi=eqReg.psi, equilibrium=eqReg.equilibrium + )[intersect_index] sfunc_orthogonal = ( lambda i: sfunc_orthogonal_original(i) @@ -1468,9 +1507,11 @@ def test_combineSfuncsPoloidalSpacingIntegrated(self, eqReg): eqReg.extend_lower = intersect_index eqReg.extend_upper = 1 - d = eqReg.totalDistance(psi=eqReg.psi) + d = eqReg.totalDistance(psi=eqReg.psi, equilibrium=eqReg.equilibrium) sfunc = eqReg.combineSfuncs(eqReg, sfunc_orthogonal) assert sfunc(0.0) == tight_approx(0.0) - assert sfunc(n - 1.0) == tight_approx(eqReg.totalDistance(psi=eqReg.psi)) + assert sfunc(n - 1.0) == tight_approx( + eqReg.totalDistance(psi=eqReg.psi, equilibrium=eqReg.equilibrium) + ) diff --git a/integrated_tests/connected_doublenull_nonorthogonal/expected_nonorthogonal.grd.nc b/integrated_tests/connected_doublenull_nonorthogonal/expected_nonorthogonal.grd.nc index a3db571b..a672f29a 100644 --- a/integrated_tests/connected_doublenull_nonorthogonal/expected_nonorthogonal.grd.nc +++ b/integrated_tests/connected_doublenull_nonorthogonal/expected_nonorthogonal.grd.nc @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:7dfb841cbbfbc89ab458c1e09e7108572521f40392bbb87164de419eb4353403 -size 638387 +oid sha256:b57f8e20d7c8155dc8d8dcf86d3613b1f247363e0a02b414f78a89a4d2df3b08 +size 671498 diff --git a/integrated_tests/connected_doublenull_orthogonal/expected_orthogonal.grd.nc b/integrated_tests/connected_doublenull_orthogonal/expected_orthogonal.grd.nc index 8f4e63bb..e862bd74 100644 --- a/integrated_tests/connected_doublenull_orthogonal/expected_orthogonal.grd.nc +++ b/integrated_tests/connected_doublenull_orthogonal/expected_orthogonal.grd.nc @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:1155390ac1606a4f89597e4a6699948b9d3aba819ebb4d597f897e71880e0d44 -size 656877 +oid sha256:e1bdde054651788e9977dd51b49a700a9d18e2b0d55a89e5d60a26159f5a501a +size 687632 diff --git a/integrated_tests/utils.py b/integrated_tests/utils.py index 7fa2ce95..7c3f3619 100644 --- a/integrated_tests/utils.py +++ b/integrated_tests/utils.py @@ -19,9 +19,12 @@ "module_versions", # Variables that have been added. These entries can be removed if/when the # expected output is re-generated. - "penalty_mask", - "closed_wall_R", - "closed_wall_Z", +] + +# Variables that are extra sensitive to rounding errors. Tolerances will be increased by +# 1e2 for these variables. +expected_large_error_vars = [ + "ShiftTorsion_ylow", ] @@ -111,11 +114,21 @@ def run_case(name, inputfile, expectedfile, *, rtol, atol, diagnose, add_noise=N .load() .drop(expected_different_vars, errors="ignore") ) + if add_noise is not None: + # Increase tolerances for some variables that are sensitive to rounding errors. + expected_large_errors = expected[expected_large_error_vars] + expected = expected.drop(expected_large_error_vars) + actual_large_errors = actual[expected_large_error_vars] + actual = actual.drop(expected_large_error_vars) if diagnose: check_errors(expected, actual, rtol=rtol, atol=atol) xrt.assert_allclose(expected, actual, rtol=rtol, atol=atol) + if add_noise is not None: + xrt.assert_allclose( + expected_large_errors, actual_large_errors, rtol=1e2 * rtol, atol=1e2 * atol + ) for attrname in expected_different_attrs: del expected.attrs[attrname]