Chunking

zoidberg/zoidberg.py

@@ -8,7 +8,7 @@
 from zoidberg import __version__
 
 from . import fieldtracer
 from .diff import field_line_length
 from .field import Slab
 from .grid import Grid
 from .poloidal_grid import StructuredPoloidalGrid
@@ -55,13 +55,7 @@
 
 
 def make_maps(
-    grid,
-    magnetic_field,
-    nslice=1,
-    quiet=False,
-    field_tracer=None,
-    refine_parallel_integral=20,
-    **kwargs,
+    grid, magnetic_field, nslice=1, quiet=False, field_tracer=None, n_chunks=1, **kwargs
 ):
     """Make the forward and backward FCI maps
 
@@ -224,65 +218,73 @@
     B_cell = [np.empty(shape) for _ in range(3)]
 
     for j in range(ny):
-        coords = None
-        y_all = None
+        chunks_x = np.array_split(np.arange(0, nx), n_chunks)
         pol, ycoord = grid.getPoloidalGrid(j)
+        num = 20
 
-        for direction in [-1, +1]:
-            # Get this poloidal grid
-            _, ycoordnext = grid.getPoloidalGrid(j + direction)
+        for chunk in chunks_x:
+            coords = None
+            y_all = None
 
-            # Get the next poloidal grid
-            pol_slice = []
-            y_slices = np.linspace(ycoord, ycoordnext, refine_parallel_integral * 2 + 1)
-            if y_all is None:
-                y_all = y_slices[::-1]
-            else:
-                y_all = np.concatenate((y_all, y_slices[1:]))
-
-            tmp = np.array(
-                field_tracer.follow_field_lines(pol.R, pol.Z, y_slices, rtol=rtol)
-            )
-            if coords is None:
-                coords = tmp[::-1]
-            else:
-                coords = np.concatenate((coords, tmp[1:]))
-                if prog is not None:
-                    prog.update()
+            for direction in [-1, +1]:
+                # Get this poloidal grid
+                _, ycoordnext = grid.getPoloidalGrid(j + direction)
 
-        for k in range(3):
-            slc = slice(
-                refine_parallel_integral * k,
-                -refine_parallel_integral * (2 - k) if k < 2 else None,
-            )
-            sg_22[k][:, j, :] = field_line_length(coords[slc], y_all[slc])
-        coords = coords[refine_parallel_integral:-refine_parallel_integral]
-        y_all = y_all[refine_parallel_integral:-refine_parallel_integral]
-
-        Bs = [
-            magnetic_field.Byfunc(coord[..., 0], coord[..., 1], y)
-            for coord, y in zip(coords, y_all)
-        ]
-
-        B_cell[0][:, j, :] = Bs[0]
-        B_cell[1][:, j, :] = Bs[refine_parallel_integral]
-        B_cell[2][:, j, :] = Bs[-1]
-
-        facs = np.ones(refine_parallel_integral * 2 + 1)
-        assert len(y_all) == len(facs)
-        facs[0] = 0.5
-        facs[-1] = 0.5
-
-        metric = pol.metric()
-        Jperp0 = np.sqrt(metric["g_xx"] * metric["g_zz"] - metric["g_xz"] ** 2)
-        B0 = Bs[len(Bs) // 2]
-        vols = [
-            Jperp0 * B0 / Bpar * fac * R
-            for Bpar, fac, R in zip(Bs, facs, coords[..., 0])
-        ]
-        assert len(vols) == len(facs)
-        J = np.sum(vols, axis=0) / (len(Bs) - 1)
-        jacobian[:, j, :] = J
+                # Get the next poloidal grid
+                pol_slice = []
+                y_slices = np.linspace(ycoord, ycoordnext, num * 2 + 1)
+                if y_all is None:
+                    y_all = y_slices[::-1]
+                else:
+                    y_all = np.concatenate((y_all, y_slices[1:]))
+
+                tmp = np.array(
+                    field_tracer.follow_field_lines(
+                        pol.R[chunk], pol.Z[chunk], y_slices, rtol=rtol
+                    )
+                )
+                if coords is None:
+                    coords = tmp[::-1]
+                else:
+                    coords = np.concatenate((coords, tmp[1:]))
+
+            for k in range(3):
+                slc = slice(num * k, -num * (2 - k) if k < 2 else None)
+                sg_22[k][chunk, j, :] = get_dist(coords[slc], y_all[slc])
+            coords = coords[num:-num]
+            y_all = y_all[num:-num]
+
+            Bs = [
+                magnetic_field.Byfunc(coord[..., 0], coord[..., 1], y)
+                for coord, y in zip(coords, y_all)
+            ]
+
+            B_cell[0][chunk, j, :] = Bs[0]
+            B_cell[1][chunk, j, :] = Bs[num]
+            B_cell[2][chunk, j, :] = Bs[-1]
+
+            facs = np.ones(num * 2 + 1)
+            assert len(y_all) == len(facs)
+            facs[0] = 0.5
+            facs[-1] = 0.5
+
+            metric = pol.metric()
+            if not isinstance(magnetic_field, Slab):
+                Jperp0 = np.sqrt(
+                    metric["g_xx"][chunk] * metric["g_zz"][chunk]
+                    - metric["g_xz"][chunk] ** 2
+                )
+            else:  # Slabs only have one metric coefficient and not the whole array stored
+                Jperp0 = np.sqrt(metric["g_xx"] * metric["g_zz"] - metric["g_xz"] ** 2)
+
+            B0 = Bs[len(Bs) // 2]
+            vols = [
+                Jperp0 * B0 / Bpar * fac * R
+                for Bpar, fac, R in zip(Bs, facs, coords[..., 0])
+            ]
+            assert len(vols) == len(facs)
+            J = np.sum(vols, axis=0) / (len(Bs) - 1)
+            jacobian[chunk, j, :] = J
 
         if prog is not None:
             prog.update()
@@ -684,6 +686,7 @@
     with MapWriter(gridfile, new_names=new_names, metric2d=metric2d, quiet=quiet) as mw:
         mw.add_grid_field(grid, magnetic_field)
         mw.add_maps(maps)
+        mw.add_dagp()
 
 
 def write_Bfield_to_vtk(