diff --git a/zoidberg/zoidberg.py b/zoidberg/zoidberg.py index 5be5141..523ee5e 100644 --- a/zoidberg/zoidberg.py +++ b/zoidberg/zoidberg.py @@ -61,6 +61,7 @@ def make_maps( quiet=False, field_tracer=None, refine_parallel_integral=20, + n_chunks=4, **kwargs, ): """Make the forward and backward FCI maps @@ -80,6 +81,8 @@ def make_maps( given field refine_parallel_integral: The number of intermediate points for g_22 calculation + n_chunks: + Number of chunks in the x-direction the calculation of J is split up. kwargs Optional arguments for field line tracing, etc. @@ -224,65 +227,79 @@ def make_maps( 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) - 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, 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[chunk], pol.Z[chunk], 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 k in range(3): + slc = slice( + refine_parallel_integral * k, + -refine_parallel_integral * (2 - k) if k < 2 else None, + ) + sg_22[k][chunk, 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][chunk, j, :] = Bs[0] + B_cell[1][chunk, j, :] = Bs[refine_parallel_integral] + B_cell[2][chunk, 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() + if isinstance(metric["g_xx"], np.ndarray) and len(metric["g_xx"].shape) > 1: + 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()