Merge pull request #33 from boutproject/more-6

More changes

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)

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)

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__ = [

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

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
 

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

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
 ):