From b7b44163d1b5ab2e40def0c9c2125c9d760d18f0 Mon Sep 17 00:00:00 2001 From: Tobias Tork Date: Mon, 13 Apr 2026 17:55:34 +0200 Subject: [PATCH 1/4] Fix numpy function and add slab exception --- zoidberg/zoidberg.py | 126 +++++++++++++++++++++++-------------------- 1 file changed, 68 insertions(+), 58 deletions(-) diff --git a/zoidberg/zoidberg.py b/zoidberg/zoidberg.py index 2d8659a..dae1047 100644 --- a/zoidberg/zoidberg.py +++ b/zoidberg/zoidberg.py @@ -54,7 +54,7 @@ def parallel_slice_field_name(field, offset): return f"{prefix}_{field}{suffix}" -def make_maps(grid, magnetic_field, nslice=1, quiet=False, field_tracer=None, **kwargs): +def make_maps(grid, magnetic_field, nslice=1, quiet=False, field_tracer=None,n_chunks = 1, **kwargs): """Make the forward and backward FCI maps Parameters @@ -211,65 +211,75 @@ def make_maps(grid, magnetic_field, nslice=1, quiet=False, field_tracer=None, ** 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) - - # 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.concat((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.concat((coords, tmp[1:])) - if prog is not None: - prog.update() - - for k in range(3): - slc = slice(num * k, -num * (2 - k) if k < 2 else None) - sg_22[k][:, 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][:, j, :] = Bs[0] - B_cell[1][:, j, :] = Bs[num] - B_cell[2][:, 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() - 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 - + + for chunk in chunks_x: + + coords = None + y_all = None + + + + + for direction in [-1, +1]: + # Get this poloidal grid + _, ycoordnext = grid.getPoloidalGrid(j + direction) + + # 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() + try: + Jperp0 = np.sqrt(metric["g_xx"][chunk] * metric["g_zz"][chunk] - metric["g_xz"][chunk] ** 2) + except: # 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() From 52ffac4cf1ba778d11f82de99142834b6e08d074 Mon Sep 17 00:00:00 2001 From: Tobias Tork Date: Mon, 13 Apr 2026 17:57:32 +0200 Subject: [PATCH 2/4] default write dagp --- zoidberg/zoidberg.py | 1 + 1 file changed, 1 insertion(+) diff --git a/zoidberg/zoidberg.py b/zoidberg/zoidberg.py index dae1047..29b7d01 100644 --- a/zoidberg/zoidberg.py +++ b/zoidberg/zoidberg.py @@ -682,6 +682,7 @@ def write_maps( 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( From de18db0fdd587d115c451e22db00848a38872b39 Mon Sep 17 00:00:00 2001 From: Tobias Tork Date: Thu, 23 Apr 2026 12:48:50 +0200 Subject: [PATCH 3/4] Formatting --- zoidberg/zoidberg.py | 73 ++++++++++++++++++++++---------------------- 1 file changed, 37 insertions(+), 36 deletions(-) diff --git a/zoidberg/zoidberg.py b/zoidberg/zoidberg.py index 29b7d01..4a223ad 100644 --- a/zoidberg/zoidberg.py +++ b/zoidberg/zoidberg.py @@ -44,9 +44,9 @@ def parallel_slice_field_name(field, offset): """ absoffset = abs(offset) if absoffset < 1: - assert ( - abs(absoffset - 0.5) < 1e-6 - ), f"Expected an offset of +- 0.5 but got {offset}" + assert abs(absoffset - 0.5) < 1e-6, ( + f"Expected an offset of +- 0.5 but got {offset}" + ) prefix = "low" if offset < 0 else "high" return f"{field}_cell_y{prefix}" prefix = "forward" if offset > 0 else "backward" @@ -54,7 +54,9 @@ def parallel_slice_field_name(field, offset): return f"{prefix}_{field}{suffix}" -def make_maps(grid, magnetic_field, nslice=1, quiet=False, field_tracer=None,n_chunks = 1, **kwargs): +def make_maps( + grid, magnetic_field, nslice=1, quiet=False, field_tracer=None, n_chunks=1, **kwargs +): """Make the forward and backward FCI maps Parameters @@ -211,22 +213,18 @@ def make_maps(grid, magnetic_field, nslice=1, quiet=False, field_tracer=None,n_c B_cell = [np.empty(shape) for _ in range(3)] for j in range(ny): - chunks_x = np.array_split(np.arange(0,nx), n_chunks) + chunks_x = np.array_split(np.arange(0, nx), n_chunks) pol, ycoord = grid.getPoloidalGrid(j) num = 20 - + for chunk in chunks_x: - coords = None y_all = None - - - - + for direction in [-1, +1]: # Get this poloidal grid _, ycoordnext = grid.getPoloidalGrid(j + direction) - + # Get the next poloidal grid pol_slice = [] y_slices = np.linspace(ycoord, ycoordnext, num * 2 + 1) @@ -234,43 +232,46 @@ def make_maps(grid, magnetic_field, nslice=1, quiet=False, field_tracer=None,n_c 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) + 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() try: - Jperp0 = np.sqrt(metric["g_xx"][chunk] * metric["g_zz"][chunk] - metric["g_xz"][chunk] ** 2) - except: # Slabs only have one metric coefficient and not the whole array stored + Jperp0 = np.sqrt( + metric["g_xx"][chunk] * metric["g_zz"][chunk] + - metric["g_xz"][chunk] ** 2 + ) + except: # 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 @@ -279,7 +280,7 @@ def make_maps(grid, magnetic_field, nslice=1, quiet=False, field_tracer=None,n_c 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() @@ -431,9 +432,9 @@ def add_dagp(self): assert self.grid, "The grid is needed to compute the DAGP. Set the grid first." assert self.is_open, "The grid file needs to be open. Call open first." - assert ( - self.field - ), "The field is needed to compute the DAGP. Set the grid first." + assert self.field, ( + "The field is needed to compute the DAGP. Set the grid first." + ) handles = {} poloidal_grids = self.grid.poloidal_grids @@ -535,9 +536,9 @@ def write_dict(self, metric, always3d=True): for k, v in metric.items(): if isinstance(v, np.ndarray): - assert np.all( - np.isfinite(v) - ), f"{k} is not finite in {v.size - np.sum(np.isfinite(v))} of {v.size} cells" + assert np.all(np.isfinite(v)), ( + f"{k} is not finite in {v.size - np.sum(np.isfinite(v))} of {v.size} cells" + ) self.f.write(k, v) def _write_par_metric(self, maps, nslice, ypar): @@ -546,9 +547,9 @@ def _write_par_metric(self, maps, nslice, ypar): meandy = np.mean(np.diff(ypar)) Ly = self.grid.Ly if self.grid else meandy * ny if self.grid and len(ypar) > 1: - assert np.isclose( - Ly, meandy * ny - ), f"Ly of grid (Ly={Ly}) does not seem to match the average dy (ny * dy = {ny} * {meandy} = {ny * meandy}" + assert np.isclose(Ly, meandy * ny), ( + f"Ly of grid (Ly={Ly}) does not seem to match the average dy (ny * dy = {ny} * {meandy} = {ny * meandy}" + ) yperiodic = self.grid.yperiodic if self.grid else True for offset in chain(range(1, nslice + 1), range(-1, -(nslice + 1), -1)): pypar = np.roll(ypar, -offset) @@ -557,9 +558,9 @@ def _write_par_metric(self, maps, nslice, ypar): if meandy * (pypar[i] - pypar[i - 1]) < 0: pypar[i] += np.sign(meandy) * Ly if len(pypar) > 1: - assert np.isclose( - np.mean(np.diff(pypar)), meandy - ), f"Mean of dy changes from {meandy} to {np.mean(np.diff(pypar))}. Values: {ypar} -> {pypar}" + assert np.isclose(np.mean(np.diff(pypar)), meandy), ( + f"Mean of dy changes from {meandy} to {np.mean(np.diff(pypar))}. Values: {ypar} -> {pypar}" + ) RZ = np.array([maps[parallel_slice_field_name(k, offset)] for k in "RZ"]) par_pgrids = [StructuredPoloidalGrid(*RZ[:, :, i]) for i in range(ny)] From fd721f1021beec82a3124424249f6c476a3b6f60 Mon Sep 17 00:00:00 2001 From: Tobias Tork Date: Thu, 23 Apr 2026 12:50:52 +0200 Subject: [PATCH 4/4] Add proper handling of slab Jperp0 --- zoidberg/zoidberg.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zoidberg/zoidberg.py b/zoidberg/zoidberg.py index 4a223ad..31adbfc 100644 --- a/zoidberg/zoidberg.py +++ b/zoidberg/zoidberg.py @@ -264,12 +264,12 @@ def make_maps( facs[-1] = 0.5 metric = pol.metric() - try: + if not isinstance(magnetic_field, Slab): Jperp0 = np.sqrt( metric["g_xx"][chunk] * metric["g_zz"][chunk] - metric["g_xz"][chunk] ** 2 ) - except: # Slabs only have one metric coefficient and not the whole array stored + 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]