Add examples from main BOUT++ repo

examples/README.md

@@ -0,0 +1,16 @@
+Example Zoidberg grid generation
+================================
+
+Here are several examples of generating grids using Zoidberg for simulations
+using the FCI parallel transform:
+
+- `example-straight-stellarator.py`: A "straight stellarator", a rotating
+  ellipse stellarator without curvature
+- `straight-stellarator-curvilinear.py`: A straight stellarator but based on
+  curvilinear grids
+- `tokamak.py`: A tokamak grid generated from an EFIT "geqdsk" file
+
+The remaining files demonstrate some of the plotting capabilities of Zoidberg:
+
+- `field_line.py`: 3D plots of magnetic field lines
+- `poincare.py`: A Poincaré plot of a straight stellarator

examples/example-straight-stellarator.py

@@ -0,0 +1,38 @@
+#!/usr/bin/env python
+
+import zoidberg
+import matplotlib.pyplot as plt
+
+nx = 5
+ny = 6
+nz = 7
+
+# Create magnetic field
+magnetic_field = zoidberg.field.StraightStellarator(radius=1.0)
+
+# Create a rectangular grid in (x,y,z)
+rectangle = zoidberg.grid.rectangular_grid(
+    nx, ny, nz, Lx=1.0, Lz=1.0, Ly=10.0, yperiodic=True
+)
+
+# Here both the field and and grid are centred at (x,z) = (0,0)
+# and the rectangular grid here fits entirely within the coils
+
+maps = zoidberg.make_maps(rectangle, magnetic_field)
+
+# Pick a poloidal slice and the next slice
+yslice = 0
+pol, ycoord = rectangle.getPoloidalGrid(yslice)
+pol_next, ycoord_next = rectangle.getPoloidalGrid(yslice + 1)
+
+# Plot the grid points at this poloidal slice
+plt.plot(pol.R, pol.Z, "x")
+
+# Get the coordinates which the forward map corresponds to
+R_next, Z_next = pol_next.getCoordinate(
+    maps["forward_xt_prime"][:, yslice, :], maps["forward_zt_prime"][:, yslice, :]
+)
+
+plt.plot(R_next, Z_next, "o")
+
+plt.show()

examples/field_line.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python
+
+# Demonstrates how to generate 3D plots of magnetic field lines
+
+import zoidberg
+
+# Size of the domain in y (periodic)
+yperiod = 10.0
+
+# Define magnetic field
+magnetic_field = zoidberg.field.StraightStellarator(
+    I_coil=0.4, radius=1.0, yperiod=yperiod
+)
+
+# Make 3D plot
+zoidberg.plot.plot_3d_field_line(
+    magnetic_field,
+    0.3,  # x starting position
+    0.0,  # z starting position
+    yperiod,  # Periodicity of y domain
+    cycles=20,
+)

examples/poincare.py

@@ -0,0 +1,23 @@
+#!/usr/bin/env python
+
+# Demonstrates how to generate Poincare plots of magnetic fields
+
+import zoidberg
+import numpy as np
+
+# Size of the domain in y (periodic)
+yperiod = 10.0
+
+# Define magnetic field
+magnetic_field = zoidberg.field.StraightStellarator(
+    I_coil=0.4, radius=1.0, yperiod=yperiod
+)
+
+# Make Poincare plot
+zoidberg.plot.plot_poincare(
+    magnetic_field,
+    np.linspace(0, 0.5, 5),  # x starting positions
+    0.0,  # z starting positions
+    yperiod,  # Periodicity of y domain
+    interactive=True,
+)  # Click on plot to add points

examples/straight-stellarator-curvilinear.py

@@ -0,0 +1,85 @@
+#!/usr/bin/env python
+
+# Create grids for a straight stellarator, based on curvilinear grids
+
+import numpy as np
+
+import zoidberg
+
+
+#############################################################################
+# Define the magnetic field
+
+# Length in y after which the coils return to their starting (R,Z) locations
+yperiod = 10.0
+
+magnetic_field = zoidberg.field.StraightStellarator(
+    I_coil=0.3, radius=1.0, yperiod=yperiod
+)
+
+#############################################################################
+# Create the inner flux surface, starting at a point at phi=0
+# To do this we need to define the y locations of the poloidal points
+# where we will construct grids
+
+start_r = 0.2
+start_z = 0.0
+
+nslices = 8  # Number of poloidal slices
+ycoords = np.linspace(0, yperiod, nslices)
+npoints = 20  # Points per poloidal slice
+
+rzcoord, ycoords = zoidberg.fieldtracer.trace_poincare(
+    magnetic_field, start_r, start_z, yperiod, y_slices=ycoords, revs=npoints
+)
+
+inner_lines = []
+for i in range(nslices):
+    r = rzcoord[:, i, 0]
+    z = rzcoord[:, i, 1]
+    line = zoidberg.rzline.line_from_points(r, z)
+    # Re-map the points so they're approximately uniform in distance along the surface
+    # Note that this results in some motion of the line
+    line = line.equallySpaced()
+    inner_lines.append(line)
+
+# Now have a list of y coordinates (ycoords) and inner lines (inner_lines)
+
+#############################################################################
+# Generate a fixed circle for the outer boundary
+
+outer_line = zoidberg.rzline.circle(R0=0.0, r=0.8)
+
+#############################################################################
+# Now have inner and outer boundaries for each poloidal grid
+# Generate a grid on each poloidal slice using the elliptic grid generator
+
+nx = 20
+ny = 20
+
+pol_grids = [
+    zoidberg.poloidal_grid.grid_elliptic(inner_line, outer_line, nx, ny)
+    for inner_line in inner_lines
+]
+
+#############################################################################
+# Create a grid, then calculate forward and backward maps
+
+grid = zoidberg.grid.Grid(pol_grids, ycoords, yperiod, yperiodic=True)
+
+maps = zoidberg.make_maps(grid, magnetic_field)
+
+#############################################################################
+# Write grid file
+
+filename = "stellarator.fci.nc"
+
+print("Writing to grid file '{0}'".format(filename))
+zoidberg.write_maps(
+    grid, magnetic_field, maps, gridfile=filename, new_names=False, metric2d=True
+)
+
+#############################################################################
+# Plot maps
+
+zoidberg.plot.plot_forward_map(grid, maps)

examples/straight-stellarator-poincare.py

@@ -0,0 +1,21 @@
+#!/usr/bin/env python
+
+# Create grids for a straight stellarator, based on
+# curvilinear grids
+
+import numpy as np
+import matplotlib.pyplot as plt
+
+import zoidberg
+
+#############################################################################
+# Define the magnetic field
+
+# Length in y after which the coils return to their starting (R,Z) locations
+yperiod = 10.0
+
+magnetic_field = zoidberg.field.StraightStellarator(
+    I_coil=0.4, radius=1.0, yperiod=yperiod
+)
+
+zoidberg.plot.plot_poincare(magnetic_field, 0.3, 0.0, yperiod)