From 31605ac8b4f4f02e686e55a7c314ba37a458b0fa Mon Sep 17 00:00:00 2001 From: Tobias Tork Date: Tue, 16 Jun 2026 09:21:39 +0200 Subject: [PATCH 1/5] Fix GEQDSK psi error --- zoidberg/field.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/zoidberg/field.py b/zoidberg/field.py index af68411..7b2f5d1 100644 --- a/zoidberg/field.py +++ b/zoidberg/field.py @@ -1237,7 +1237,8 @@ def __init__(self, gfile): print("Height: {0} -> {1} m".format(self.zmin, self.zmax)) # Poloidal flux - self.psi = np.transpose(data["psi"]) + #self.psi = np.transpose(data["psi"]) + self.psi = data["psi"] nr, nz = self.psi.shape # Normalising factors: psi on axis and boundary From b1b0b446bb64eef931fb4872253cce6f13fd16a2 Mon Sep 17 00:00:00 2001 From: Tobias Tork Date: Tue, 16 Jun 2026 09:23:50 +0200 Subject: [PATCH 2/5] Use full field for area calculation --- zoidberg/zoidberg.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zoidberg/zoidberg.py b/zoidberg/zoidberg.py index b99963a..7a9b498 100644 --- a/zoidberg/zoidberg.py +++ b/zoidberg/zoidberg.py @@ -269,7 +269,7 @@ def make_maps( y_all = y_all[refine_parallel_integral:-refine_parallel_integral] Bs = [ - magnetic_field.Byfunc(coord[..., 0], coord[..., 1], y) + magnetic_field.Bmag(coord[..., 0], coord[..., 1], y) for coord, y in zip(coords, y_all) ] From e509c513073f8b70396b56f948e2cb904cd13436 Mon Sep 17 00:00:00 2001 From: Tobias Tork Date: Tue, 16 Jun 2026 09:24:12 +0200 Subject: [PATCH 3/5] Add dagp to default mapwriter --- zoidberg/zoidberg.py | 1 + 1 file changed, 1 insertion(+) diff --git a/zoidberg/zoidberg.py b/zoidberg/zoidberg.py index 7a9b498..fbe2888 100644 --- a/zoidberg/zoidberg.py +++ b/zoidberg/zoidberg.py @@ -700,6 +700,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 d30b26ce068b2c030b3e435656072cbcfbb1f20e Mon Sep 17 00:00:00 2001 From: Tobias Tork Date: Tue, 16 Jun 2026 09:28:49 +0200 Subject: [PATCH 4/5] Ruff formatting --- zoidberg/field.py | 2 +- zoidberg/zoidberg.py | 114 ++++++++++++++++++++++++++++--------------- 2 files changed, 75 insertions(+), 41 deletions(-) diff --git a/zoidberg/field.py b/zoidberg/field.py index 7b2f5d1..e1a9cd0 100644 --- a/zoidberg/field.py +++ b/zoidberg/field.py @@ -1237,7 +1237,7 @@ def __init__(self, gfile): print("Height: {0} -> {1} m".format(self.zmin, self.zmax)) # Poloidal flux - #self.psi = np.transpose(data["psi"]) + # self.psi = np.transpose(data["psi"]) self.psi = data["psi"] nr, nz = self.psi.shape diff --git a/zoidberg/zoidberg.py b/zoidberg/zoidberg.py index fbe2888..cdf521c 100644 --- a/zoidberg/zoidberg.py +++ b/zoidberg/zoidberg.py @@ -54,6 +54,11 @@ def parallel_slice_field_name(field, offset): return f"{prefix}_{field}{suffix}" +import gc +import psutil +import os + + def make_maps( grid, magnetic_field, @@ -161,55 +166,84 @@ def make_maps( prog = tqdm(total=total_work, desc="Tracing") else: update_progress(0, **kwargs) - for slice_index, parallel_slices in enumerate(parallel_slices_list): - for j in range(ny): - # Get this poloidal grid - pol, ycoord = grid.getPoloidalGrid(j) - - # Get the next poloidal grid - pol_slices = [] - y_slices = [ycoord] - for parallel_slice in parallel_slices: - pol_slice, y_slice = grid.getPoloidalGrid(j + parallel_slice.offset) - pol_slices.append(pol_slice) - y_slices.append(y_slice) - - # We only want the end point, as [0,...] is the initial position - coords = field_tracer.follow_field_lines(pol.R, pol.Z, y_slices, rtol=rtol)[ - 1:, ... - ] - for parallel_slice, coord, pol_slice in zip( - parallel_slices, coords, pol_slices - ): - # Store the coordinates in real space - parallel_slice.R[:, j, :] = coord[:, :, 0] - parallel_slice.Z[:, j, :] = coord[:, :, 1] - - # Get the indices into the slice poloidal grid - if pol_slice is None: - # No slice grid, so hit a boundary - xind = -1 - zind = -1 - else: - # Find the indices for these new locations on the slice poloidal grid - xcoord = coord[:, :, 0] - zcoord = coord[:, :, 1] - xind, zind = pol_slice.findIndex(xcoord, zcoord) + process = psutil.Process(os.getpid()) + + prev = process.memory_info().rss / 1e6 - # Check boundary defined by the field - outside = magnetic_field.boundary.outside(xcoord, y_slice, zcoord) - xind[outside] = -1 - zind[outside] = -1 + for slice_index, parallel_slices in enumerate(parallel_slices_list): + chunks_x = np.array_split(np.arange(0, nx), n_chunks) - parallel_slice.xt_prime[:, j, :] = xind - parallel_slice.zt_prime[:, j, :] = zind + for j in range(ny): + print("Difference: ", process.memory_info().rss / 1e6 - prev, "MB") + prev = np.copy(process.memory_info().rss / 1e6) + for chunk in chunks_x: + # print("\nStart ny:", process.memory_info().rss / 1e6, "MB\n") + # print("Difference: ", process.memory_info().rss / 1e6 - prev, "MB" ) + prev = np.copy(process.memory_info().rss / 1e6) + # Get this poloidal grid + pol, ycoord = grid.getPoloidalGrid(j) + + # Get the next poloidal grid + pol_slices = [] + y_slices = [ycoord] + for parallel_slice in parallel_slices: + pol_slice, y_slice = grid.getPoloidalGrid(j + parallel_slice.offset) + pol_slices.append(pol_slice) + y_slices.append(y_slice) + + # We only want the end point, as [0,...] is the initial position + coords = None + + # before = len(gc.get_objects()) + # print("Before tracing:", process.memory_info().rss / 1e6, "MB\n") + coords = field_tracer.follow_field_lines( + pol.R[chunk], pol.Z[chunk], y_slices, rtol=rtol + )[1:, ...] + # print("After tracing:", process.memory_info().rss / 1e6, "MB\n") + + # after = len(gc.get_objects()) + # print(after - before) + + for parallel_slice, coord, pol_slice in zip( + parallel_slices, coords, pol_slices + ): + # Store the coordinates in real space + + parallel_slice.R[chunk, j, :] = coord[:, :, 0].copy() + parallel_slice.Z[chunk, j, :] = coord[:, :, 1].copy() + + # Get the indices into the slice poloidal grid + if pol_slice is None: + # No slice grid, so hit a boundary + xind = -1 + zind = -1 + else: + # Find the indices for these new locations on the slice poloidal grid + xcoord = coord[:, :, 0].copy() + zcoord = coord[:, :, 1].copy() + + xind, zind = pol_slice.findIndex(xcoord, zcoord) + + # # Check boundary defined by the field + # outside = magnetic_field.boundary.outside(xcoord, y_slice, zcoord) + # xind[outside] = -1 + # zind[outside] = -1 + + parallel_slice.xt_prime[chunk, j, :] = xind + parallel_slice.zt_prime[chunk, j, :] = zind + + # print("Before delete:", process.memory_info().rss / 1e6, "MB\n") + del coords, xcoord, zcoord, xind, zind, coord, pol_slices + gc.collect() + # print("After delete:", process.memory_info().rss / 1e6, "MB\n") if (not quiet) and (ny > 1): if tqdm: prog.update() else: update_progress((slice_index * ny + j + 1) / total_work, **kwargs) + for k in list(maps.keys()): if "_cell_y" in k and "t_prime" in k: del maps[k] From d2bf94755c3d2fb2491a679809da32984b31b693 Mon Sep 17 00:00:00 2001 From: Tobias Tork Date: Tue, 16 Jun 2026 09:31:03 +0200 Subject: [PATCH 5/5] Revert "Ruff formatting" This reverts commit d30b26ce068b2c030b3e435656072cbcfbb1f20e. --- zoidberg/field.py | 2 +- zoidberg/zoidberg.py | 114 +++++++++++++++---------------------------- 2 files changed, 41 insertions(+), 75 deletions(-) diff --git a/zoidberg/field.py b/zoidberg/field.py index e1a9cd0..7b2f5d1 100644 --- a/zoidberg/field.py +++ b/zoidberg/field.py @@ -1237,7 +1237,7 @@ def __init__(self, gfile): print("Height: {0} -> {1} m".format(self.zmin, self.zmax)) # Poloidal flux - # self.psi = np.transpose(data["psi"]) + #self.psi = np.transpose(data["psi"]) self.psi = data["psi"] nr, nz = self.psi.shape diff --git a/zoidberg/zoidberg.py b/zoidberg/zoidberg.py index cdf521c..fbe2888 100644 --- a/zoidberg/zoidberg.py +++ b/zoidberg/zoidberg.py @@ -54,11 +54,6 @@ def parallel_slice_field_name(field, offset): return f"{prefix}_{field}{suffix}" -import gc -import psutil -import os - - def make_maps( grid, magnetic_field, @@ -166,84 +161,55 @@ def make_maps( prog = tqdm(total=total_work, desc="Tracing") else: update_progress(0, **kwargs) - - process = psutil.Process(os.getpid()) - - prev = process.memory_info().rss / 1e6 - for slice_index, parallel_slices in enumerate(parallel_slices_list): - chunks_x = np.array_split(np.arange(0, nx), n_chunks) - for j in range(ny): - print("Difference: ", process.memory_info().rss / 1e6 - prev, "MB") - prev = np.copy(process.memory_info().rss / 1e6) - for chunk in chunks_x: - # print("\nStart ny:", process.memory_info().rss / 1e6, "MB\n") - # print("Difference: ", process.memory_info().rss / 1e6 - prev, "MB" ) - prev = np.copy(process.memory_info().rss / 1e6) - # Get this poloidal grid - pol, ycoord = grid.getPoloidalGrid(j) + # Get this poloidal grid + pol, ycoord = grid.getPoloidalGrid(j) + + # Get the next poloidal grid + pol_slices = [] + y_slices = [ycoord] + for parallel_slice in parallel_slices: + pol_slice, y_slice = grid.getPoloidalGrid(j + parallel_slice.offset) + pol_slices.append(pol_slice) + y_slices.append(y_slice) + + # We only want the end point, as [0,...] is the initial position + coords = field_tracer.follow_field_lines(pol.R, pol.Z, y_slices, rtol=rtol)[ + 1:, ... + ] - # Get the next poloidal grid - pol_slices = [] - y_slices = [ycoord] - for parallel_slice in parallel_slices: - pol_slice, y_slice = grid.getPoloidalGrid(j + parallel_slice.offset) - pol_slices.append(pol_slice) - y_slices.append(y_slice) - - # We only want the end point, as [0,...] is the initial position - coords = None - - # before = len(gc.get_objects()) - # print("Before tracing:", process.memory_info().rss / 1e6, "MB\n") - coords = field_tracer.follow_field_lines( - pol.R[chunk], pol.Z[chunk], y_slices, rtol=rtol - )[1:, ...] - # print("After tracing:", process.memory_info().rss / 1e6, "MB\n") - - # after = len(gc.get_objects()) - # print(after - before) - - for parallel_slice, coord, pol_slice in zip( - parallel_slices, coords, pol_slices - ): - # Store the coordinates in real space - - parallel_slice.R[chunk, j, :] = coord[:, :, 0].copy() - parallel_slice.Z[chunk, j, :] = coord[:, :, 1].copy() - - # Get the indices into the slice poloidal grid - if pol_slice is None: - # No slice grid, so hit a boundary - xind = -1 - zind = -1 - else: - # Find the indices for these new locations on the slice poloidal grid - xcoord = coord[:, :, 0].copy() - zcoord = coord[:, :, 1].copy() - - xind, zind = pol_slice.findIndex(xcoord, zcoord) - - # # Check boundary defined by the field - # outside = magnetic_field.boundary.outside(xcoord, y_slice, zcoord) - # xind[outside] = -1 - # zind[outside] = -1 - - parallel_slice.xt_prime[chunk, j, :] = xind - parallel_slice.zt_prime[chunk, j, :] = zind - - # print("Before delete:", process.memory_info().rss / 1e6, "MB\n") - del coords, xcoord, zcoord, xind, zind, coord, pol_slices - gc.collect() - # print("After delete:", process.memory_info().rss / 1e6, "MB\n") + for parallel_slice, coord, pol_slice in zip( + parallel_slices, coords, pol_slices + ): + # Store the coordinates in real space + parallel_slice.R[:, j, :] = coord[:, :, 0] + parallel_slice.Z[:, j, :] = coord[:, :, 1] + + # Get the indices into the slice poloidal grid + if pol_slice is None: + # No slice grid, so hit a boundary + xind = -1 + zind = -1 + else: + # Find the indices for these new locations on the slice poloidal grid + xcoord = coord[:, :, 0] + zcoord = coord[:, :, 1] + xind, zind = pol_slice.findIndex(xcoord, zcoord) + + # Check boundary defined by the field + outside = magnetic_field.boundary.outside(xcoord, y_slice, zcoord) + xind[outside] = -1 + zind[outside] = -1 + + parallel_slice.xt_prime[:, j, :] = xind + parallel_slice.zt_prime[:, j, :] = zind if (not quiet) and (ny > 1): if tqdm: prog.update() else: update_progress((slice_index * ny + j + 1) / total_work, **kwargs) - for k in list(maps.keys()): if "_cell_y" in k and "t_prime" in k: del maps[k]