diff --git a/examples/rotating-ellipse.py b/examples/rotating-ellipse.py index 6bf8e06..2482d58 100755 --- a/examples/rotating-ellipse.py +++ b/examples/rotating-ellipse.py @@ -1,11 +1,34 @@ #!/usr/bin/env python -# Create grids for a rotating ellipse stellarator, based on curvilinear grids - +import argparse import numpy as np - import zoidberg +parser = argparse.ArgumentParser( + prog="RotatingEllipse", + description="Create grids for a rotating ellipse stellarator, based on curvilinear grids.", +) +parser.add_argument( + "-nx", type=int, default=20, help="Number of radial (x) cells including boundaries" +) +parser.add_argument( + "-ny", "--nslices", type=int, default=8, help="Number of parallel (y) slices" +) +parser.add_argument("-nz", type=int, default=20, help="Number of poloidal (z) cells") +parser.add_argument( + "-ns", + "--nsamples-per-dim", + type=int, + default=1, + help="Number of samples per dimension", +) +parser.add_argument("-pb", "--partial-boundaries", action="store_true", default=True) +parser.add_argument("-o", "--output", type=str, default="") +args = parser.parse_args() + +if args.output == "": + args.output = f"rotating-ellipse-{args.nx}x{args.nslices}x{args.nz}-s{args.nsamples_per_dim}{'-npb' if not args.partial_boundaries else ''}.fci.nc" + ############################################################################# # Define the magnetic field @@ -30,8 +53,7 @@ start_r = xcentre + 0.4 start_z = 0.0 -nslices = 8 # Number of poloidal slices -ycoords = np.linspace(0, yperiod, nslices) +ycoords = np.linspace(0, yperiod, args.nslices, endpoint=False) npoints = 20 # Points per poloidal slice rzcoord, ycoords = zoidberg.fieldtracer.trace_poincare( @@ -39,7 +61,7 @@ ) inner_lines = [] -for i in range(nslices): +for i in range(args.nslices): r = rzcoord[:, i, 0, 0] z = rzcoord[:, i, 0, 1] line = zoidberg.rzline.line_from_points(r, z) @@ -59,11 +81,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 -nz = 20 - pol_grids = [ - zoidberg.poloidal_grid.grid_elliptic(inner_line, outer_line, nx, nz) + zoidberg.poloidal_grid.grid_elliptic(inner_line, outer_line, args.nx, args.nz) for inner_line in inner_lines ] @@ -72,19 +91,14 @@ grid = zoidberg.grid.Grid(pol_grids, ycoords, yperiod, yperiodic=True) -samples_per_dim = 2 # Sub-sampling in each cell -partial_boundaries = True +maps = zoidberg.make_maps(grid, magnetic_field, samples_per_dim=args.nsamples_per_dim) -maps = zoidberg.make_maps(grid, magnetic_field, samples_per_dim=samples_per_dim) - -maps = zoidberg.weights.modify_maps(maps, partial_boundaries=partial_boundaries) +maps = zoidberg.weights.modify_maps(maps, partial_boundaries=args.partial_boundaries) ############################################################################# # Write grid file -filename = f"rotating-ellipse-{nx}x{nslices}x{nz}-s{samples_per_dim}{'-npb' if not partial_boundaries else ''}.fci.nc" - -print(f"Writing to grid file '{filename}'") +print(f"Writing to grid file '{args.output}'") zoidberg.write_maps( - grid, magnetic_field, maps, gridfile=filename, new_names=False, metric2d=False + grid, magnetic_field, maps, gridfile=args.output, new_names=False, metric2d=False ) diff --git a/zoidberg/grid.py b/zoidberg/grid.py index ea384eb..947f612 100644 --- a/zoidberg/grid.py +++ b/zoidberg/grid.py @@ -49,9 +49,27 @@ def __init__(self, poloidal_grids, ycoords, Ly, yperiodic=False, name="fci_grid" self.poloidal_grids = poloidal_grids self.ycoords = np.asarray(ycoords) + + if len(self.ycoords.shape) != 1: + raise ValueError("ycoords must be 1-dimensional") + + def is_monotonic(a): + return np.all(a[1:] >= a[:-1]) + + if not is_monotonic(self.ycoords): + raise ValueError("ycoords must be monotonically increasing") + self.Ly = Ly self.yperiodic = yperiodic + if yperiodic: + if ( + Ly <= self.ycoords[-1] - self.ycoords[0] + 1e-10 + ): # Note: Catch case where they are equal + raise ValueError( + f"Ly={Ly} must be > max(ycoords) - min(ycoords) = {ycoords[-1]} - {ycoords[0]}" + ) + self._metric_cache = None # Define the shape of the grid