diff --git a/test/mms/delp_grid.py b/test/mms/delp_grid.py index 1e7b96a..699c9d1 100644 --- a/test/mms/delp_grid.py +++ b/test/mms/delp_grid.py @@ -12,6 +12,11 @@ ] +class FixedCurvedSlab(zb.field.CurvedSlab): + def Rfunc(self, x, z, phi): + return x + + def gen_name(*args): nx, ny, nz, R0, r0, r1, mode = args return f"delp_{modes[mode][0]}_{nx}_{ny}_{nz}_{R0}.fci.nc" @@ -27,7 +32,7 @@ def gen_grid(nx, ny, nz, R0, r0, r1, mode=0): Z = np.sin(theta) * r pol_grid = zb.poloidal_grid.StructuredPoloidalGrid(R, Z) - field = zb.field.CurvedSlab(Bz=0, Bzprime=0, Rmaj=R0) + field = FixedCurvedSlab(Bz=0, Bzprime=0, Rmaj=R0) grid = zb.grid.Grid(pol_grid, phi, 5, yperiodic=True) fn = gen_name(*args) diff --git a/test/mms/laplace_grid.py b/test/mms/laplace_grid.py index bfdd6d2..d53f5c3 100644 --- a/test/mms/laplace_grid.py +++ b/test/mms/laplace_grid.py @@ -15,6 +15,11 @@ ] +class FixedCurvedSlab(zb.field.CurvedSlab): + def Rfunc(self, x, z, phi): + return x + + def gen_name(*args): nx, ny, nz, R0, r0, r1, mode = args return f"laplace_{modes[mode][0]}_{nx}_{ny}_{nz}_{R0}.fci.nc" @@ -40,7 +45,7 @@ def gen_grid(nx, ny, nz, R0, r0, r1, mode=0): plt.show() pol_grid = zb.poloidal_grid.StructuredPoloidalGrid(R, Z) - field = zb.field.CurvedSlab(Bz=0, Bzprime=0, Rmaj=R0) + field = FixedCurvedSlab(Bz=0, Bzprime=0, Rmaj=R0) grid = zb.grid.Grid(pol_grid, phi, 5, yperiodic=True) fn = gen_name(*args) diff --git a/zoidberg/__init__.py b/zoidberg/__init__.py index fe3dfcf..88884c8 100644 --- a/zoidberg/__init__.py +++ b/zoidberg/__init__.py @@ -9,9 +9,7 @@ __version__ = get_version(root="..", relative_to=__file__) -import zoidberg - -from . import field, fieldtracer, grid, plot +from . import field, fieldtracer, grid, plot, zoidberg from .zoidberg import make_maps, write_maps __all__ = [ diff --git a/zoidberg/field.py b/zoidberg/field.py index 3e75617..2cad2fc 100644 --- a/zoidberg/field.py +++ b/zoidberg/field.py @@ -1545,7 +1545,7 @@ def field_values(cls, r, phi, z, configuration=0, plot_poincare=False): print("Making poincare plot (only works on IPP network)...") ## Create configuration objects config = tracer.types.MagneticConfig() - config.configIds = configuration + config = _set_config(config, configuration) pos = tracer.types.Points3D() pos.x1 = np.linspace(5.6, 6.2, 80) @@ -1673,7 +1673,7 @@ def magnetic_axis(self, phi_axis=0, configuration=0): ) config = tracer.types.MagneticConfig() - config.configIds = configuration + config = _set_config(config, configuration) settings = tracer.types.AxisSettings() res = tracer.service.findAxis(0.05, config, settings) @@ -1742,7 +1742,7 @@ def _calc_chunk(cls, x, configuration): ## Create configuration objects config = tracer.types.MagneticConfig() - config.configIds = configuration + config = _set_config(config, configuration) ## Call tracer service redo = 2 @@ -1766,6 +1766,57 @@ def _calc_chunk(cls, x, configuration): return np.array((res.field.x1, res.field.x2, res.field.x3)).reshape(x.shape) +def _set_config(config, configuration): + if isinstance(configuration, int): + if configuration >= 0: + config.configIds = configuration + return config + known = { + -1: [1, 1, 1, 1, 1, 1, 1], + "FMM002": [13423, 13423, 13423, 13423, 13423, -3544, -3544], + } + currents = known[configuration] + currents = np.array([108] * 5 + [36] * 2) * np.array(currents) + print(currents) + coils = ( + [160, 165, 170, 175, 180, 185, 190, 195, 200, 205] + + [161, 166, 171, 176, 181, 186, 191, 196, 201, 206] + + [162, 167, 172, 177, 182, 187, 192, 197, 202, 207] + + [163, 168, 173, 178, 183, 188, 193, 198, 203, 208] + + [164, 169, 174, 179, 184, 189, 194, 199, 204, 209] + + [210, 212, 214, 216, 218, 220, 222, 224, 226, 228] + + [211, 213, 215, 217, 219, 221, 223, 225, 227, 229] + + [230, 231, 232, 233, 234, 235, 236, 237, 238, 239] + + [350, 241, 351, 352, 353] + ) + + if len(currents) == 7: + currents = np.array([[x] * 10 for x in currents]).flatten() + elif len(currents) == 7 + 10 + 5: + currents = np.append( + np.array([[x] * 10 for x in currents[:7]]).flatten(), currents[7:] + ) + elif len(currents) == 10: + currents = np.append( + np.array([[x] * 10 for x in currents[:7]]).flatten(), + np.append(np.array([currents[7:9]] * 5), [currents[9]] * 5), + ) + if len(currents) == 70: + coils = coils[:70] + # I1, I1, I1, I1, I1, I1, I1, I1, I1, I1, + # I2, I2, I2, I2, I2, I2, I2, I2, I2, I2, + # I3, I3, I3, I3, I3, I3, I3, I3, I3, I3, + # I4, I4, I4, I4, I4, I4, I4, I4, I4, I4, + # I5, I5, I5, I5, I5, I5, I5, I5, I5, I5, + # IA, IA, IA, IA, IA, IA, IA, IA, IA, IA, + # IB, IB, IB, IB, IB, IB, IB, IB, IB, IB, + # Icc1, Icc2, Icc3, Icc4, Icc5, Icc6, Icc7, Icc8, Icc9, Icc10, + # Itc1, Itc2, Itc3, Itc4, Itc5] + config.coilsIds = coils + config.coilsIdsCurrents = np.array(currents) + return config + + class W7X_vacuum_on_demand: def __init__(self, configuration): self.configuration = configuration diff --git a/zoidberg/fieldtracer.py b/zoidberg/fieldtracer.py index a7c809d..34877cf 100644 --- a/zoidberg/fieldtracer.py +++ b/zoidberg/fieldtracer.py @@ -3,6 +3,8 @@ import numpy as np from scipy.integrate import odeint +from .field import _set_config + class FieldTracer(object): """A class for following magnetic field lines @@ -607,7 +609,7 @@ def getConfig(self): self.config.inverseField = False return self.config config = self.flt.types.MagneticConfig() - config.configIds = self.configId + config = _set_config(config, self.configId) config.inverseField = False return config diff --git a/zoidberg/poloidal_grid.py b/zoidberg/poloidal_grid.py index 611c086..2114c00 100644 --- a/zoidberg/poloidal_grid.py +++ b/zoidberg/poloidal_grid.py @@ -522,7 +522,7 @@ def metric(self): "g_xz": g[..., 0, 1], "gzz": ginv[..., 1, 1], "g_zz": g[..., 1, 1], - "J": JB, + # "J": JB, } @@ -603,8 +603,10 @@ def grid_elliptic( return_coords=False, nx_outer=0, nx_inner=0, + legacy_align=False, inner_ort=True, maxfac_inner=None, + dz_relax=None, ): """Create a structured grid between inner and outer boundaries using elliptic method @@ -715,17 +717,21 @@ def grid_elliptic( longer = inner ind = np.argmax(shorter.R) shorter = rzline.RZline(np.roll(shorter.R, -ind), np.roll(shorter.Z, -ind)) - dr = shorter.R - longer.R[:, None] - dz = shorter.Z - longer.Z[:, None] - delta = dr**2 + dz**2 - fac = len(longer.R) / len(shorter.R) - j = np.arange(len(shorter.R), dtype=int) - sums = [ - np.sum(delta[np.round(j * fac).astype(int) - i, j]) - for i in range(len(longer.R)) - ] - ind = -np.argmin(sums) - longer = rzline.RZline(np.roll(longer.R, -ind), np.roll(longer.Z, -ind)) + if legacy_align: + ind = np.argmax(longer.R) + longer = rzline.RZline(np.roll(longer.R, -ind), np.roll(longer.Z, -ind)) + else: + dr = shorter.R - longer.R[:, None] + dz = shorter.Z - longer.Z[:, None] + delta = dr**2 + dz**2 + fac = len(longer.R) / len(shorter.R) + j = np.arange(len(shorter.R), dtype=int) + sums = [ + np.sum(delta[np.round(j * fac).astype(int) - i, j]) + for i in range(len(longer.R)) + ] + ind = -np.argmin(sums) + longer = rzline.RZline(np.roll(longer.R, -ind), np.roll(longer.Z, -ind)) if len(inner.R) < len(outer.R): inner = shorter outer = longer @@ -821,6 +827,7 @@ def getfac(x): return_coords=True, inner_ort=inner_ort, maxfac_inner=maxfac_inner, + dz_relax=dz_relax, ) # Note: Lower case x,z are indices @@ -918,11 +925,11 @@ def getfac(x): Z_zm = Z_zm[1:-1, :] Z_zp = Z_zp[1:-1, :] - eps = -0.7 - dRdz = (R_zp - R_zm) / (2.0 * dz * (1 + eps)) + dz_relax = dz_relax or 10 / 3 + dRdz = dz_relax * (R_zp - R_zm) / (2.0 * dz) dRdx = (R_xp - R_xm) / (2.0 * dx) - dZdz = (Z_zp - Z_zm) / (2.0 * dz * (1 + eps)) + dZdz = dz_relax * (Z_zp - Z_zm) / (2.0 * dz) dZdx = (Z_xp - Z_xm) / (2.0 * dx) a = dRdz**2 + dZdz**2 diff --git a/zoidberg/rzline.py b/zoidberg/rzline.py index f01bdba..cc8b1b3 100644 --- a/zoidberg/rzline.py +++ b/zoidberg/rzline.py @@ -189,9 +189,8 @@ def positionPolygon(self, theta=None): if theta is None: return self.R, self.Z n = len(self.R) - theta = np.remainder(theta, 2.0 * pi) dtheta = 2.0 * np.pi / n - ind = np.trunc(theta / dtheta).astype(int) + ind = np.trunc(theta / dtheta).astype(int) % n rem = np.remainder(theta, dtheta) indp = (ind + 1) % n return (rem * self.R[indp] + (1.0 - rem) * self.R[ind]), ( @@ -462,7 +461,7 @@ def line_from_points_poly(rarray, zarray, show=False, spline_order=None): rl, zl = line.position(theta) - ind = np.floor(float(i) * theta / (2.0 * np.pi)) + ind = int(np.floor(float(i) * theta / (2.0 * np.pi))) # Insert after this index @@ -485,6 +484,107 @@ def line_from_points_poly(rarray, zarray, show=False, spline_order=None): return RZline(rvals, zvals, spline_order=spline_order) +def line_from_points_fast(rarray, zarray, **kw): + """ + Find a periodic line which goes through the given (r,z) points. + + This version is particularly fast, but fails if the points are not + sufficiently close to a circle. + """ + rz = rarray, zarray + cent = [np.mean(x) for x in rz] + dist = [a - b for (a, b) in zip(rz, cent)] + angle = np.arctan2(*dist) + ind = np.argsort(angle) + return RZline(*[x[ind] for x in rz], **kw) + + +def line_from_points_two_opt(rarray, zarray, opt=1e-3, **kwargs): + """This is probably way to slow. + + Use the two opt algorithm: + From https://stackoverflow.com/a/44080908 + License: CC-BY-SA 4.0 + """ + # Calculate the euclidian distance in n-space of the route r traversing cities c, ending at the path start. + path_distance = lambda r, c: np.sum( + [np.linalg.norm(c[r[p]] - c[r[p - 1]]) for p in range(len(r))] + ) + # Reverse the order of all elements from element i to element k in array r. + two_opt_swap = lambda r, i, k: np.concatenate( + (r[0:i], r[k : -len(r) + i - 1 : -1], r[k + 1 : len(r)]) + ) + + # 2-opt Algorithm adapted from https://en.wikipedia.org/wiki/2-opt + def two_opt(cities, improvement_threshold): + # Make an array of row numbers corresponding to cities. + route = np.arange(cities.shape[0]) + improvement_factor = 1 + best_distance = path_distance(route, cities) + while improvement_factor > improvement_threshold: + # Record the distance at the beginning of the loop. + distance_to_beat = best_distance + # From each city except the first and last, + for swap_first in range(1, len(route) - 2): + # to each of the cities following, + for swap_last in range(swap_first + 1, len(route)): + # try reversing the order of these cities + new_route = two_opt_swap(route, swap_first, swap_last) + # and check the total distance with this modification. + new_distance = path_distance(new_route, cities) + # If the path distance is an improvement, + if new_distance < best_distance: + route = new_route + best_distance = new_distance + # Calculate how much the route has improved. + improvement_factor = 1 - best_distance / distance_to_beat + return route + + l = line_from_points_fast(rarray, zarray, **kwargs) + cities = np.array([l.R, l.Z]).T + route = two_opt(cities, opt) + R, Z = cities[route].T + return RZline(R, Z, **kwargs) + + +def line_from_points_convex_hull(rarray, zarray, **kwargs): + """ + Find a periodic line which goes through the given (r,z) points + + This function uses an alphashape to find the boundary, and then projects + the points on the boundary for sorting. + + This requires shapely + """ + import shapely as sg + + rz = np.array((rarray, zarray)).T + shape = sg.MultiPoint(rz).convex_hull.exterior + theta = [shape.project(sg.Point(*p)) for p in rz] + ind = np.argsort(theta) + return RZline(*rz[ind].T) + + +def line_from_points_alphashape(rarray, zarray, alpha=2.0, **kwargs): + """ + Find a periodic line which goes through the given (r,z) points + + This function uses an alphashape to find the boundary, and then projects + the points on the boundary for sorting. + + This requires shapely and alphashape + """ + import alphashape + import shapely as sg + + rz = np.array((rarray, zarray)).T + shape = alphashape.alphashape(rz, alpha) + shape = shape.exterior + theta = [shape.project(sg.Point(*p)) for p in rz] + ind = np.argsort(theta) + return RZline(*rz[ind].T) + + def line_from_points( rarray, zarray, show=False, spline_order=None, is_sorted=False, smooth=False ): diff --git a/zoidberg/smooth.py b/zoidberg/smooth.py index 187df7f..1758f06 100644 --- a/zoidberg/smooth.py +++ b/zoidberg/smooth.py @@ -1,5 +1,6 @@ import numpy as np from scipy import linalg +from scipy.optimize import least_squares def find_outliers(r, z, cutoff=0.75, bounds=10): @@ -51,9 +52,11 @@ def rollup(original_array, startindex, length, replacement): return original_array -def gen_newr(r, z, splitedlist, bounds): +def gen_newr(r, z, splitedlist, bounds, plot): newr = r.copy() newz = z.copy() + have_failed = False + deb = [] for alist in splitedlist: assert len(alist) >= bounds if len(alist) >= len(r) / 2: @@ -68,45 +71,234 @@ def gen_newr(r, z, splitedlist, bounds): seconds = 0 diff_plus_r = r[end_index] - r[seconds] diff_plus_z = z[end_index] - z[seconds] - A = np.array( - [ - [1, 0.5, 1 / 300, 0, 0, 0], - [0, 0, 0, 1, 0.5, 1 / 300], - [diff_minus_z, 0, 0, -diff_minus_r, 0, 0], + fac = 1 + if 0: + A = np.array( [ - diff_plus_z, - diff_plus_z, - diff_plus_z / 100, - -diff_plus_r, - -diff_plus_r, - -diff_plus_r / 100, - ], - ] - ) + [1, 0.5, 1 / 3 / fac, 1 / 4, 0, 0, 0, 0], + [0, 0, 0, 0, 1, 0.5, 1 / 3 / fac, 1 / 4], + [diff_minus_z, 0, 0, 0, -diff_minus_r, 0, 0, 0], + [ + diff_plus_z, + diff_plus_z, + diff_plus_z / fac, + diff_plus_z, + -diff_plus_r, + -diff_plus_r, + -diff_plus_r / fac, + -diff_plus_r, + ], + ] + ) + else: + A = np.array( + [ + [1, 0.5, 1 / 3 / fac, 0, 0, 0], + [0, 0, 0, 1, 0.5, 1 / 3 / fac], + [diff_minus_z, 0, 0, -diff_minus_r, 0, 0], + [ + diff_plus_z, + diff_plus_z, + diff_plus_z / fac, + -diff_plus_r, + -diff_plus_r, + -diff_plus_r / fac, + ], + ] + ) + b = np.array( - [[r[end_index] - r[start_index]], [z[end_index] - z[start_index]], [0], [0]] + [r[end_index] - r[start_index], z[end_index] - z[start_index], 0, 0] ) - x = linalg.lstsq(A, b) + + def lenofs(x): + if len(x) == 6: + x1r, x2r, x3r = x[:3] + x1z, x2z, x3z = x[3:] + xr = x[:3] + xz = x[3:] + xs = xr * xr + xz * xz + x1s, x2s, x3s = xs + x12 = x1r * x2r + x1z * x2z + x13 = x1r * x3r + x1z * x3z + x23 = x2r * x3r + x2z * x3z + num = 100 + s = np.linspace(0, 1, num) + tmp = np.sqrt( + x1s + + x2s * s**2 + + x3s * s**4 + + 2 * (x12 * s + x13 * s**2 + x23 * s**3) + ) + return np.sum(tmp) / num + x1r, x2r, x3r, x4r = x[:4] + x1z, x2z, x3z, x4z = x[4:] + xr = x[:4] + xz = x[4:] + xs = xr * xr + xz * xz + x1s, x2s, x3s, x4s = xs + x12 = x1r * x2r + x1z * x2z + x13 = x1r * x3r + x1z * x3z + x14 = x1r * x4r + x1z * x4z + x23 = x2r * x3r + x2z * x3z + x24 = x2r * x4r + x2z * x4z + x34 = x3r * x4r + x3z * x4z + # return x1s + x12 + 2 / 3 * x13 + 1/2*x14 +1/3 * x23 + 1 / 4 * x2s + 1 / 5 * x3s +1/7*x4s + num = 100 + s = np.linspace(0, 1, num) + tmp = np.sqrt( + x1s + + x2s * s**2 + + x3s * s**4 + + x4s * s**6 + + 2 + * ( + x12 * s + + x13 * s**2 + + x14 * s**3 + + x23 * s**3 + + x24 * s**4 + + x34 * s**5 + ) + ) + return np.sum(tmp) / num + + def mysin(a, b): + aa, ab = a + ba, bb = b + return (aa * bb - ab * ba) / np.sqrt( + (aa * aa + ab * ab + 1e-8) * (ba * ba + bb * bb + 1e-8) + ) + + def mycos(a, b): + aa, ab = a + ba, bb = b + return (aa * ba + ab * bb) / np.sqrt( + (aa * aa + ab * ab + 1e-8) * (ba * ba + bb * bb + 1e-8) + ) + + def curv(x, f): + k = len(x) // 2 + a = np.sum(x[:k] * np.arange(k)) + b = x[1] + c = np.sum(x[k:] * np.arange(k)) + d = x[k + 1] + return (a * a + c * c) * f, (b * b + d * d) * f + + def fun(x, A, b, signs): + ret = A @ x - b + k = len(x) // 2 + newsigns = (x[0], x[1], np.sum(x[:k]), np.sum(x[k:])) + mycoss = mycos(signs[:2], newsigns[:2]), mycos(signs[2:], newsigns[2:]) + return ( + *ret, + lenofs(x) / 1e5, + # (1 - mycoss[0]) / 1e2, + # (mycoss[1] - 1) / 1e2, + *(curv(x, 1 / 1e3)), + ) + + if 0: + x = linalg.lstsq(A, b) + x0 = x[0] + else: + signs = (diff_minus_r, diff_minus_z, -diff_plus_r, -diff_plus_z) + x = least_squares( + fun, + np.zeros(A.shape[1]), + args=(A, b, signs), + gtol=None, + xtol=1e-14, + ftol=None, + ) + # if np.max(np.abs(x.fun)) > 1e-5: + if x.optimality > 1e-8: + have_failed = True + deb += [[x.cost, x.optimality, np.max(np.abs(x.fun))]] + continue + x0 = x.x + newsigns = (x0[0], x0[1], np.sum(x0[:3]), np.sum(x0[3:])) + S = np.linspace(0, 1, num=len(alist)) - X = ( - r[start_index] - + x[0][0] * S - + (x[0][1] * S**2) * 0.5 - + (x[0][2] * S**3) / 300 - ) - Y = ( - z[start_index] - + x[0][3] * S - + (x[0][4] * S**2) * 0.5 - + (x[0][5] * S**3) / 300 - ) + + if len(x0) == 6: + X = ( + r[start_index] + + x0[0] * S + + (x0[1] * S**2) * 0.5 + + (x0[2] * S**3) / 3 / fac + ) + Y = ( + z[start_index] + + x0[3] * S + + (x0[4] * S**2) * 0.5 + + (x0[5] * S**3) / 3 / fac + ) + else: + X = ( + r[start_index] + + x0[0] * S + + (x0[1] * S**2) * 0.5 + + (x0[2] * S**3) / 3 / fac + + x0[3] * S**4 / 4 + ) + Y = ( + z[start_index] + + x0[4] * S + + (x0[5] * S**2) * 0.5 + + (x0[6] * S**3) / 3 / fac + + x0[7] * S**4 / 4 + ) + + # Check spacing at beginning and end + dstart = (X[0] - newr[start_index - 1]) ** 2 + ( + Y[0] - newz[start_index - 1] + ) ** 2 + dend = (X[-1] - newr[start_index + len(alist)]) ** 2 + ( + Y[-1] - newz[start_index + len(alist)] + ) ** 2 + maxfac = 3 + if not (maxfac**-2 < dstart / dend < maxfac**2): + print(f"dstart / dend is {dstart/dend} - check has failed") + deb += [[dstart, dend, 3]] + have_failed = True + if start_index < 0: newr = rollup(newr, start_index, len(alist), X) newz = rollup(newz, start_index, len(alist), Y) else: newr[start_index : start_index + len(alist)] = X newz[start_index : start_index + len(alist)] = Y - return newr, newz + if plot: + from matplotlib.collections import LineCollection + + S = np.linspace(0, 1, 1000) + + X = ( + r[start_index] + + x0[0] * S + + (x0[1] * S**2) * 0.5 + + (x0[2] * S**3) / 3 / fac + ) + Y = ( + z[start_index] + + x0[3] * S + + (x0[4] * S**2) * 0.5 + + (x0[5] * S**3) / 3 / fac + ) + points = np.array([X, Y]).T.reshape(-1, 1, 2) + segments = np.concatenate([points[:-1], points[1:]], axis=1) + norm = plt.Normalize(S.min(), S.max()) + lc = LineCollection(segments, cmap="viridis", norm=norm) + lc.set_array(S) + lc.set_linewidth(2) + line = plt.gca().add_collection(lc) + # plt.plot(X, Y, "k-") + + return newr, newz, have_failed, deb + + +alldeb = [] def smooth(r, z, cutoff=0.75, bounds=10, bounds_find=None, plot=False): @@ -131,20 +323,48 @@ def smooth(r, z, cutoff=0.75, bounds=10, bounds_find=None, plot=False): if plot: import matplotlib.pyplot as plt + # print(r, z) plt.plot(r, z, "o-", label="input") - outliers = find_outliers(r, z, cutoff, bounds_find or bounds) - if plot: - plt.plot(*np.array([(r[i], z[i]) for i in outliers]).T, "x", label="outlier") - splitedlist = generate_points_tosmooth(bounds, outliers, len(r)) - if plot: - # plt.plot( - print(splitedlist) - ret = gen_newr(r, z, splitedlist, bounds) + + bounds_find = bounds_find or bounds + bounds0 = bounds + debs = [] + while True: + outliers = [] + for b in range(1, bounds_find + 1): + # print(find_outliers(r, z, cutoff, b)) + outliers += find_outliers(r, z, cutoff, b) + if plot: + # print(np.array([(r[i], z[i]) for i in outliers]).T) + if outliers: + plt.plot( + *np.array([(r[i], z[i]) for i in outliers]).T, "x", label="outlier" + ) + splitedlist = generate_points_tosmooth(bounds, outliers, len(r)) + if plot: + # plt.plot( + # print(splitedlist) + pass + a, b, fail, deb = gen_newr(r, z, splitedlist, bounds, plot) + if not fail: + ret = a, b + break + r, z = a, b + bounds += 1 + debs += [deb[-1]] + if bounds > bounds0 * 5: + plt.legend() + plt.gca().set_aspect("equal") + plt.title("failed") + alldeb.append(debs) + return + + alldeb.append(debs) if plot: - plt.plot(*ret, label="new") + plt.plot(*ret, "o-", label="new") plt.legend() plt.gca().set_aspect("equal") - plt.show() + # plt.show() return ret @@ -158,12 +378,24 @@ def smooth(r, z, cutoff=0.75, bounds=10, bounds_find=None, plot=False): fn = sys.argv[1] dat = np.loadtxt(fn) dat.shape = (-1, 2, dat.shape[1]) - line = zb.rzline.RZline(*dat[1]) - r, z = dat[1] - r = [] - z = [] - - newr, newz = smooth(r, z, cutoff=0.75, bounds=10) - plt.plot(newr, newz, "x-") - plt.plot(r, z) + # line = zb.rzline.RZline(*dat[1]) + j = 0 + for i in range(len(dat)): + r, z = dat[i] + if find_outliers(r, z, 0.75, 10): + plt.figure() + smooth(r, z, cutoff=0.75, bounds=10, plot=True) + j += 1 + print(f"figure {j}") + # if j == 2: + # break + + if any([len(x) for x in alldeb]): + f, axs = plt.subplots(3, 1) + for deb in alldeb: + x = np.array(deb) + print(deb) + print(x.shape) + for a, b in zip(axs, x.T): + a.plot(b) plt.show() diff --git a/zoidberg/stencil_dagp_fv.py b/zoidberg/stencil_dagp_fv.py new file mode 100755 index 0000000..cbdb183 --- /dev/null +++ b/zoidberg/stencil_dagp_fv.py @@ -0,0 +1,361 @@ +#!/usr/bin/env python +# coding: utf-8 + +import time + +time_start = time.time() # noqa + +import sys + +import numpy as np +from boututils.datafile import DataFile as DF + +verbose = 1 + + +def toCent(RZ): + RZc = RZ + np.roll(RZ, 1, axis=-1) + RZc = RZc[:, 1:] + RZc[:, :-1] + return RZc / 4 + + +def l2(x): + return np.sqrt(np.sum(x**2, axis=0)) + + +org_print = print + + +def my_print(*args): + args = f"{time.time() - time_start:10.3f} s: ", *args + org_print(*args) + + +def log(*args): + if verbose > 1: + my_print(*args) + + +def print(*args): + if verbose: + my_print(*args) + + +def load(fn): + with DF(fn) as f: + RZ = [f[k] for k in "RZ"] + log("opened") + RZ = np.array(RZ) + _, a, b, c = RZ.shape + + coefsX = np.zeros((a - 1, b, c, 2)) + coefsZ = np.zeros((a - 2, b, c, 2)) + for d, x1 in zip((coefsX, coefsZ), "XZ"): + for i, x2 in enumerate("XZ"): + key = f"dagp_fv_{x1}{x2}" + fixup2(d[..., i], f[key]) + key = "dagp_fv_volume" + volume = f[key] + log("done reading") + return RZ, volume, coefsX, coefsZ + + +def doit(pols, plot=False): + RZs = [] + + ### Calculate Volume of the cell + # + ### Go in a line around the cell + + A = np.empty((pols[0].nx, len(pols), pols[0].nz)) + Ar = np.empty_like(A) + + for gi, g in enumerate(pols): + n = 50 + x0 = np.arange(g.nx)[1:-1, None, None] + y0 = np.arange(g.nz)[None, :, None] + x1 = x0[:, :, 0] + y1 = y0[:, :, 0] + x = [] + y = [] + zero = np.zeros((g.nx - 2, g.nz, n)) + x += [np.linspace(x1 - 0.5, x1 + 0.5, n).transpose(1, 2, 0)] + y += [y0 + 0.5] + x += [x0 + 0.5] + y += [np.linspace(y1 + 0.5, y1 - 0.5, n).transpose(1, 2, 0)] + x += [np.linspace(x1 + 0.5, x1 - 0.5, n).transpose(1, 2, 0)] + y += [y0 - 0.5] + x += [x0 - 0.5] + y += [np.linspace(y1 - 0.5, y1 + 0.5, n).transpose(1, 2, 0)] + x = np.concatenate([k + zero for k in x], axis=-1) + y = np.concatenate([k + zero for k in y], axis=-1) + RZ = np.array(g.getCoordinate(x, y)) + + dy = RZ[1, ..., 1:] - RZ[1, ..., :-1] + xx = (RZ[0, ..., :-1] + RZ[0, ..., 1:]) / 2 + + A[1:-1, gi] = -np.sum(xx * dy, axis=-1) + Ar[1:-1, gi] = -np.sum(0.5 * xx * xx * dy, axis=-1) + + RZs += [RZ] + RZs = np.array(RZs) + print(RZs.shape) + # 3.160 s: (4, 2, 6, 32, 200) + RZs = RZs.transpose(1, 2, 0, 3, 4) + print(RZs.shape) + + volume = Ar + # f = + # V = [f[k] for k in ("dx", "dy", "dz", "J")] + + RZ = np.array( + [ + g.getCoordinate( + *np.meshgrid(np.arange(g.nx), np.arange(g.nz), indexing="ij") + ) + for g in pols + ] + ).transpose(1, 2, 0, 3) + + RZ = np.array([(g.R, g.Z) for g in pols]).transpose(1, 2, 0, 3) + + # volume = dx * dy * dz * J + # volume = V[0] * V[1] * V[2] * V[3] + # A.shape, Ar.shape, RZ.shape + + # AreaXplus + # Z + # Λ + # | a b c + # | + # | ------------ + # | | X + # | | X + # | h | o X d + # | | X + # | | X + # | ------------ + # | + # | g f e + # | + # |------------------------> x + # + # We are calculating the area of the surface denoted by X + # Note that we need to split it in two parts. + + log("calculating pos ...") + # pos = RZs[..., :n] + tmp = np.meshgrid( + np.arange(g.nx - 1) + 0.5, + np.arange(g.nz) - 0.5, + np.linspace(0, 1, n), + indexing="ij", + ) + tmp[1] += tmp[2] + print(tmp[0]) + print(tmp[1]) + pos = np.array([g.getCoordinate(*tmp[:2]) for g in pols]).transpose(1, 2, 0, 3, 4) + log("done") + # direction of edge, length ~ 1/nx + dR = pos[..., :-1] - pos[..., 1:] + dX = np.sqrt(np.sum(dR**2, axis=0)) + area = dX * (pos[0, ..., 1:] + pos[0, ..., :-1]) * 0.5 + assert area.shape[0] > 2 + # Vector in normal direction + assert dR.shape[0] == 2 + dRr = np.array((-dR[1], dR[0])) + dRr /= np.sqrt(np.sum(dRr**2, axis=0)) + dRr *= area + dRr = np.sum(dRr, axis=-1) + print(np.nanmean(dRr)) + + # vector in derivative direction + dxR = RZ[:, 1:] - RZ[:, :-1] + print("Maximum radial grid distance:", np.max(l2(dxR))) + dzR = np.roll(RZ, -1, axis=-1) - np.roll(RZ, 1, axis=-1) + dzR = 0.5 * (dzR[:, 1:] + dzR[:, :-1]) + + # dxR /= (np.sum(dxR**2, axis=0)) + # dzR /= (np.sum(dzR**2, axis=0)) + log("starting solve") + dxzR = np.array((dxR, dzR)).transpose(2, 3, 4, 1, 0) + coefsX = np.linalg.solve(dxzR, dRr.transpose(1, 2, 3, 0)) + log("done") + + # AreaZplus + # Z + # Λ + # | a b c + # | + # | -XXXXXXXXXX- + # | | | + # | | | + # | h | o | d + # | | | + # | | | + # | ------------ + # | + # | g f e + # | + # |------------------------> x + + log("concatenate") + tmp = np.meshgrid( + np.arange(g.nx - 2) + 0.5, + np.arange(g.nz) + 0.5, + np.linspace(0, 1, n), + indexing="ij", + ) + tmp[0] += tmp[2] + pos = np.array([g.getCoordinate(*tmp[:2]) for g in pols]).transpose(1, 2, 0, 3, 4) + + dR = pos[..., :-1] - pos[..., 1:] + dX = np.sqrt(np.sum(dR**2, axis=0)) + area = dX * (pos[0, ..., 1:] + pos[0, ..., :-1]) * 0.5 + log("get normal vector") + # Vector in normal direction + dRr = np.array((dR[1], -dR[0])) + dRr /= np.sqrt(np.sum(dRr**2, axis=0)) + dRr *= area + dRr = np.sum(dRr, axis=-1) + # vector in derivative direction + dxR = RZ[:, 2:] - RZ[:, :-2] + dxR = 0.5 * (np.roll(dxR, -1, axis=-1) + dxR) + dzR = (np.roll(RZ, -1, axis=-1) - RZ)[:, 1:-1] + dxzR = np.array((dxR, dzR)).transpose(2, 3, 4, 1, 0) + log("solving again") + coefsZ = np.linalg.solve(dxzR, dRr.transpose(1, 2, 3, 0)) + log("done") + + test(RZ, volume, coefsX, coefsZ, plot=plot) + + return write(RZ, volume, coefsX, coefsZ) + + +def test(RZ, volume, coefsX, coefsZ, plot=False): + inp = np.sin(RZ[1]) + ana = -np.sin(RZ[1]) + if 0: + inp = RZ[1] ** 2 + ana = RZ[1] ** 0 + if 1: + inp = np.sin(RZ[0]) + ana = -np.sin(RZ[0]) + np.cos(RZ[0]) / RZ[0] + if 0: + inp = RZ[0] ** 3 + ana = 6 * np.sin(RZ[0]) + + def xp(ijk): + return (ijk[0] + 1, *ijk[1:]) + + def xm(ijk): + return (ijk[0] - 1, *ijk[1:]) + + def zp(ijk): + return _zp(*ijk) + + def _zp(i, j, k): + return (i, j, (k + 1) % zmax) + + def zm(i, j, k): + return (i, j, (k - 1) % zmax) + + log("testing") + result = np.zeros(inp.shape) + results = [np.zeros(inp.shape) for _ in range(4)] + # dx = + dz2 = np.roll(inp, -1, axis=-1) - np.roll(inp, 1, axis=-1) + i, j, k = inp.shape + zmax = k + dx = inp[1:] - inp[:-1] + dz = 0.5 * (dz2[1:] + dz2[:-1]) + this = -(coefsX[..., 0] * dx + coefsX[..., 1] * dz) + result[:-1] -= this + result[1:] += this + for r, t in zip(results, (-coefsX[..., 0] * dx, -coefsX[..., 1] * dz)): + r[:-1] -= t + r[1:] += t + print("expect 0:", np.max(np.abs((result - results[0] - results[1])))) + if 1: + dx2 = inp[2:] - inp[:-2] + dx2 = 0.5 * (np.roll(dx2, -1, axis=-1) + dx2) + dz2 = (np.roll(inp, -1, axis=-1) - inp)[1:-1] + t1 = coefsZ[..., 0] * dx2 + t2 = coefsZ[..., 1] * dz2 + this = -(t1 + t2) + result[1:-1] -= this + result[1:-1] += np.roll(this, 1, -1) + for r, t in zip(results[2:], (-t1, -t2)): + r[1:-1] -= t + # np.roll(r, -1, -1)[1:-1] += t + r[1:-1] += np.roll(t, 1, -1) + + result[0] = 0 + result[-1] = 0 + for r in results: + r[0] = 0 + r[-1] = 0 + + if plot: + import matplotlib.pyplot as plt + + f, axs = plt.subplots(1, 3, sharex=True, sharey=True) + res = result / volume + for d, ax, t in zip( + (res, ana, res - ana), + axs, + ("computed", "analytical", "error"), + ): + slc = slice(1, -1), 1 + per = np.percentile(d[slc], [0, 2, 98, 100]) + print(per) + p = ax.pcolormesh(*[k[slc] for k in RZ], d[slc], vmin=per[1], vmax=per[2]) + ax.set_title(t) + plt.colorbar(p, ax=ax) + plt.show() + print("error:", np.mean(l2(result[1:-1] / volume[1:-1] - ana[1:-1]))) + + +def fixup(RZ, d): + c1 = np.zeros(RZ[0].shape) + if c1.shape[0] == d.shape[0] + 1: + c1[:-1] = d + else: + c1[1:-1] = d + # c1 = np.roll(c1, -1, -1) + return c1 + + +def fixup2(d, val): + if val.shape[0] == d.shape[0] + 1: + d[:] = val[:-1] + else: + d[:] = val[1:-1] + + +def write(RZ, volume, coefsX, coefsZ): + ret = {} + log("writing") + for d, x1 in zip((coefsX, coefsZ), "XZ"): + for i, x2 in enumerate("XZ"): + key = f"dagp_fv_{x1}{x2}" + ret[key] = fixup(RZ, d[..., i]) + key = "dagp_fv_volume" + ret[key] = volume + log("done") + print(ret.keys()) + return ret + + +if __name__ == "__main__": + for flag, step in (("-v", +1), ("-q", -1)): + while flag in sys.argv: + verbose += step + sys.argv.remove(flag) + plot = "-p" in sys.argv + if plot: + sys.argv.remove("-p") + + for fn in sys.argv[1:]: + print(fn) + test(*load(fn), plot=plot) diff --git a/zoidberg/zoidberg.py b/zoidberg/zoidberg.py index e4f0d24..3675039 100644 --- a/zoidberg/zoidberg.py +++ b/zoidberg/zoidberg.py @@ -8,6 +8,8 @@ from zoidberg import __version__ from . import fieldtracer +from .grid import Grid +from .poloidal_grid import StructuredPoloidalGrid try: from tqdm.auto import tqdm @@ -177,6 +179,57 @@ def make_maps(grid, magnetic_field, nslice=1, quiet=False, field_tracer=None, ** return maps +def update_metric_names(metric): + # Translate between output variable names and metric names + # Map from new to old names. Anything not in this dict + # is unchanged + name_changes = { + "g_yy": "g_22", + "gyy": "g22", + "gxx": "g11", + "gxz": "g13", + "gzz": "g33", + "g_xx": "g_11", + "g_xz": "g_13", + "g_zz": "g_33", + } + return {name_changes.get(key, key): value for key, value in metric.items()} + + +def get_metric(grid, magnetic_field): + nx, ny, nz = grid.shape + # Get metric tensor + metric = grid.metric() + + # Check if the magnetic field is in cylindrical coordinates + # If so, we need to change the gyy and g_yy metrics + pol_grid, ypos = grid.getPoloidalGrid(0) + Rmaj = magnetic_field.Rfunc(pol_grid.R, pol_grid.Z, ypos) + if Rmaj is not None: + # In cylindrical coordinates + Rmaj = np.zeros(grid.shape) + for yindex in range(grid.numberOfPoloidalGrids()): + pol_grid, ypos = grid.getPoloidalGrid(yindex) + Rmaj[:, yindex, :] = magnetic_field.Rfunc(pol_grid.R, pol_grid.Z, ypos) + metric["gyy"] = 1.0 / Rmaj**2 + metric["g_yy"] = Rmaj**2 + + # Get magnetic field and pressure + Bmag = np.zeros(grid.shape) + pressure = np.zeros(grid.shape) + print("starting Bfield stuff") + for yindex in range(grid.numberOfPoloidalGrids()): + pol_grid, ypos = grid.getPoloidalGrid(yindex) + Bmag[:, yindex, :] = magnetic_field.Bmag(pol_grid.R, pol_grid.Z, ypos) + pressure[:, yindex, :] = magnetic_field.pressure(pol_grid.R, pol_grid.Z, ypos) + By = magnetic_field.Byfunc(pol_grid.R, pol_grid.Z, ypos) + metric["g_yy"][:, yindex, :] *= (Bmag[:, yindex, :] / By) ** 2 + metric["gyy"][:, yindex, :] *= (By / Bmag[:, yindex, :]) ** 2 + print("done Bfield stuff") + + return metric, Bmag, pressure + + def write_maps( grid, magnetic_field, @@ -215,40 +268,36 @@ def write_maps( """ + metric, Bmag, pressure = get_metric(grid, magnetic_field) nx, ny, nz = grid.shape - # Get metric tensor - metric = grid.metric() - # Check if the magnetic field is in cylindrical coordinates - # If so, we need to change the gyy and g_yy metrics - pol_grid, ypos = grid.getPoloidalGrid(0) - Rmaj = magnetic_field.Rfunc(pol_grid.R, pol_grid.Z, ypos) - if Rmaj is not None: - # In cylindrical coordinates - Rmaj = np.zeros(grid.shape) - for yindex in range(grid.numberOfPoloidalGrids()): - pol_grid, ypos = grid.getPoloidalGrid(yindex) - Rmaj[:, yindex, :] = magnetic_field.Rfunc(pol_grid.R, pol_grid.Z, ypos) - metric["gyy"] = 1.0 / Rmaj**2 - metric["g_yy"] = Rmaj**2 + # Add Rxy, Bxy + metric["Rxy"] = maps["R"] + metric["Bxy"] = Bmag - # Get magnetic field and pressure - Bmag = np.zeros(grid.shape) - pressure = np.zeros(grid.shape) - print("starting Bfield stuff") - for yindex in range(grid.numberOfPoloidalGrids()): - pol_grid, ypos = grid.getPoloidalGrid(yindex) - Bmag[:, yindex, :] = magnetic_field.Bmag(pol_grid.R, pol_grid.Z, ypos) - pressure[:, yindex, :] = magnetic_field.pressure(pol_grid.R, pol_grid.Z, ypos) - By = magnetic_field.Byfunc(pol_grid.R, pol_grid.Z, ypos) - metric["g_yy"][:, yindex, :] = ( - metric["g_yy"][:, yindex, :] * (Bmag[:, yindex, :] / By) ** 2 - ) - metric["gyy"][:, yindex, :] = ( - metric["gyy"][:, yindex, :] * (By / Bmag[:, yindex, :]) ** 2 + nslice = maps["MYG"] + # Loop over offsets {1, ... nslice, -1, ... -nslice} + for offset in chain(range(1, nslice + 1), range(-1, -(nslice + 1), -1)): + par_pgrids, ypar = np.array( + [grid.getPoloidalGrid(i + offset) for i in range(ny)], dtype=object + ).T + par_pgrids = [StructuredPoloidalGrid(p.R, p.Z) for p in par_pgrids] + par_grid = Grid( + par_pgrids, + ypar, + Ly=grid.Ly, + yperiodic=grid.yperiodic, ) - print("done Bfield stuff") + + par_metric, _, _ = get_metric(par_grid, magnetic_field) + if not new_names: + par_metric = update_metric_names(par_metric) + + for k, v in par_metric.items(): + name = parallel_slice_field_name(k, offset) + assert name not in metric + metric[name] = v # Get attributes from magnetic field (e.g. psi) attributes = {} @@ -273,12 +322,9 @@ def write_maps( pass # Make dz a constant metric["dz"] = metric["dz"][0, 0] - # Add Rxy, Bxy - metric["Rxy"] = maps["R"][:, :, 0] - metric["Bxy"] = Bmag[:, :, 0] - else: - metric["Rxy"] = maps["R"] - metric["Bxy"] = Bmag + + if not new_names: + metric = update_metric_names(metric) with bdata.DataFile(gridfile, write=True, create=True, format=format) as f: f.write_file_attribute("title", "BOUT++ grid file") @@ -301,32 +347,8 @@ def write_maps( f.write("ixseps2", ixseps) # Metric tensor - - if new_names: - for key, val in metric.items(): - f.write(key, val) - else: - # Translate between output variable names and metric names - # Map from new to old names. Anything not in this dict - # is output unchanged - name_changes = { - "g_yy": "g_22", - "gyy": "g22", - "gxx": "g11", - "gxz": "g13", - "gzz": "g33", - "g_xx": "g_11", - "g_xz": "g_13", - "g_zz": "g_33", - } - for key in metric: - name = key - if name in name_changes: - name = name_changes[name] - f.write(name, metric[key]) - - if "J" in metric: - f.write("J", metric["J"]) + for key in metric: + f.write(key, metric[key]) # Magnetic field f.write("B", Bmag)