From e16b72f4cf88d734f132234abd961708bfd85463 Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Fri, 30 May 2025 15:20:28 -0700 Subject: [PATCH 01/11] weights: Calculate interpolation weights for forward/backward Create matrix representation of interpolation, saved in CSR format in the BOUT++ grid file. --- zoidberg/__init__.py | 3 +- zoidberg/weights.py | 220 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 222 insertions(+), 1 deletion(-) create mode 100644 zoidberg/weights.py diff --git a/zoidberg/__init__.py b/zoidberg/__init__.py index 88884c8..5d41ee0 100644 --- a/zoidberg/__init__.py +++ b/zoidberg/__init__.py @@ -9,7 +9,7 @@ __version__ = get_version(root="..", relative_to=__file__) -from . import field, fieldtracer, grid, plot, zoidberg +from . import field, fieldtracer, grid, plot, zoidberg, weights from .zoidberg import make_maps, write_maps __all__ = [ @@ -19,6 +19,7 @@ "grid", "make_maps", "plot", + "weights", "write_maps", "zoidberg", ] diff --git a/zoidberg/weights.py b/zoidberg/weights.py new file mode 100644 index 0000000..291c3c9 --- /dev/null +++ b/zoidberg/weights.py @@ -0,0 +1,220 @@ +"""Routines to calculate interpolation weights + +N cells divided between Ne evolving cells and Nb = N - Ne boundary cells + + [0 .. evolving cells .. (Ne - 1) | Ne .. boundary cells .. (N - 1)] + +Boundary cells include both radial (X) boundaries and parallel (Yup/Ydown) cells. + +Cell index numbers are stored in three arrays: + + cell_number[x,y,z] <- These can be calculated, not stored + cell_number_yup[x,y,z] + cell_number_ydown[x,y,z] + + +For forward and backward maps the Nw weights are stored in CSR format + + weights[Nw] + column_index[Nw] + row_index[Ne] <- Starting index into weights and column_index + This will be -1 for boundary points + Needs to be considered when getting the weights + +For each evolving cell i in 0...(Ne-1) the weight index j is +row_index[i]..(row_index[i+1] - 1) + +i.e. + +result[i] = sum_{j = row_index[i]}^{row_index[i+1]-1} weight[j] * input[column_index[j]] + +The column_index cells go from 0..(N-1), including boundary cells. +Note: row_index[i+1] may be -1, so skip over -1 entries. + +Note that these operators are usually represented as non-square +matrices: Input (columns) of length N, output (rows) of length Ne < N. +This is because boundary conditions are set independently. + + +The output grid file will contain + int cell_number(x, y, z) ; + int total_cells ; + + int forward_cell_number(x, y, z) ; + double forward_weights(t) ; + int forward_columns(t) ; + int forward_rows(t2) ; + + int backward_cell_number(x, y, z) ; + double backward_weights(t3) ; + int backward_columns(t3) ; + int backward_rows(t2) ; + +""" +import numpy as np + + +def calc_cell_numbers(maps): + """ Given a field line map dictionary, assign numbers to evolving + cells and boundary cells """ + nx, ny, nz = maps["R"].shape + MXG = maps["MXG"] + # Number of evolving cells + N_evolving = (nx - 2 * MXG) * ny * nz + + # Numbering + cell_number_array = np.zeros((nx, ny, nz), dtype=int) + cell_number = 0 + for i in range(MXG, nx - MXG): + for j in range(ny): + for k in range(nz): + cell_number_array[i, j, k] = cell_number + cell_number += 1 + + # Inner radial boundary cells + for i in range(MXG): + for j in range(ny): + for k in range(nz): + cell_number_array[i, j, k] = cell_number + cell_number += 1 + # Outer radial boundary cells + for i in range(nx - MXG, nx): + for j in range(ny): + for k in range(nz): + cell_number_array[i, j, k] = cell_number + cell_number += 1 + + backward_cell_number = np.zeros((nx, ny, nz), dtype=int) + forward_cell_number = np.zeros((nx, ny, nz), dtype=int) + + # Iterate through for forward and backward maps + forward_xt_prime = maps["forward_xt_prime"] + backward_xt_prime = maps["backward_xt_prime"] + + # Number of radial boundary cells + N_radial = 2 * MXG * ny * nz + cell_number = N_evolving + N_radial + + # Add the parallel boundary cells + for i in range(MXG, nx - MXG): + for j in range(ny): + for k in range(nz): + if backward_xt_prime[i, j, k] < 0.0: + backward_cell_number[i, j, k] = cell_number + cell_number += 1 + elif backward_xt_prime[i, j, k] >= nx: + backward_cell_number[i, j, k] = cell_number + cell_number += 1 + if forward_xt_prime[i, j, k] < 0.0: + forward_cell_number[i, j, k] = cell_number + cell_number += 1 + elif forward_xt_prime[i, j, k] >= nx: + forward_cell_number[i, j, k] = cell_number + cell_number += 1 + + return { + "N_cells": cell_number, + "N_evolving": N_evolving, + "cell_number": cell_number_array, + "forward_cell_number": forward_cell_number, + "backward_cell_number": backward_cell_number, + } + + +def calc_interpolation(cell_number, MXG, yoffset, xtarr, ztarr): + """ + Calculate CSR format matrix representing a 2D (X-Z) interpolation operation + + Implements the cubic Catmull-Rom spline + Coefficients taken from https://en.wikipedia.org/wiki/Cubic_Hermite_spline + """ + + # Offsets and weights for 1D interpolation + offsets1D = [-1, 0, 1, 2] + + def weights1D(u): + return 0.5 * np.array( + [ + -u ** 3 + 2.0 * u ** 2 - u, + 3.0 * u ** 3 - 5.0 * u ** 2 + 2, + -3.0 * u ** 3 + 4.0 * u ** 2 + u, + u ** 3 - u ** 2, + ] + ) + + # CSR format + weights = [] + columns = [] + rows = [] + + nx, ny, nz = cell_number.shape + + weight_number = 0 # Track location in weights & columns arrays + for i in range(MXG, nx - MXG): + for j in range(ny): + for k in range(nz): + xt = xtarr[i, j, k] + zt = ztarr[i, j, k] + if (xt < 0.0) or (xt >= nx): + # Boundary + rows.append(-1) + else: + # Not a boundary point => Interpolating + rows.append(weight_number) + xi = int(xt) # Floor + zi = int(zt) + + weights_x = weights1D(xt - xi) + weights_z = weights1D(zt - zi) + + for xo, xw in zip(offsets1D, weights_x): + for zo, zw in zip(offsets1D, weights_z): + columns.append( + cell_number[ + np.clip(xi + xo, 0, nx - 1), + (j + yoffset + ny) % ny, + (zi + zo + nz) % nz, + ] + ) + weights.append(xw * zw) + weight_number += 1 + return {"weights": weights, "columns": columns, "rows": rows} + + +def calc_weights(maps): + """ + Calculate interpolation weights for forward and backward maps. + + Returns a dictionary of arrays to be read into BOUT++ + """ + + numbering = calc_cell_numbers(maps) + + forward = calc_interpolation( + numbering["cell_number"], + maps["MXG"], + +1, + maps["forward_xt_prime"], + maps["forward_zt_prime"], + ) + + backward = calc_interpolation( + numbering["cell_number"], + maps["MXG"], + -1, + maps["backward_xt_prime"], + maps["backward_zt_prime"], + ) + + return { + "cell_number": numbering["cell_number"], + "total_cells": numbering["N_cells"], + "forward_cell_number": numbering["forward_cell_number"], + "forward_weights": forward["weights"], + "forward_columns": forward["columns"], + "forward_rows": forward["rows"], + "backward_cell_number": numbering["backward_cell_number"], + "backward_weights": backward["weights"], + "backward_columns": backward["columns"], + "backward_rows": backward["rows"], + } From 534efc0d6fa69ac2837239886ba9c592520ce7be Mon Sep 17 00:00:00 2001 From: bendudson Date: Fri, 30 May 2025 22:22:02 +0000 Subject: [PATCH 02/11] Apply black/isort changes --- zoidberg/__init__.py | 2 +- zoidberg/boundary.py | 4 +--- zoidberg/weights.py | 35 ++++++++++++++++++----------------- 3 files changed, 20 insertions(+), 21 deletions(-) diff --git a/zoidberg/__init__.py b/zoidberg/__init__.py index 5d41ee0..bb35a34 100644 --- a/zoidberg/__init__.py +++ b/zoidberg/__init__.py @@ -9,7 +9,7 @@ __version__ = get_version(root="..", relative_to=__file__) -from . import field, fieldtracer, grid, plot, zoidberg, weights +from . import field, fieldtracer, grid, plot, weights, zoidberg from .zoidberg import make_maps, write_maps __all__ = [ diff --git a/zoidberg/boundary.py b/zoidberg/boundary.py index 9dc48ca..406e5e0 100644 --- a/zoidberg/boundary.py +++ b/zoidberg/boundary.py @@ -1,6 +1,4 @@ -"""Boundary objects that define an 'outside' - -""" +"""Boundary objects that define an 'outside'""" import numpy as np diff --git a/zoidberg/weights.py b/zoidberg/weights.py index 291c3c9..c7c8ac1 100644 --- a/zoidberg/weights.py +++ b/zoidberg/weights.py @@ -22,7 +22,7 @@ Needs to be considered when getting the weights For each evolving cell i in 0...(Ne-1) the weight index j is -row_index[i]..(row_index[i+1] - 1) +row_index[i]..(row_index[i+1] - 1) i.e. @@ -38,25 +38,26 @@ The output grid file will contain int cell_number(x, y, z) ; - int total_cells ; + int total_cells ; - int forward_cell_number(x, y, z) ; - double forward_weights(t) ; - int forward_columns(t) ; - int forward_rows(t2) ; + int forward_cell_number(x, y, z) ; + double forward_weights(t) ; + int forward_columns(t) ; + int forward_rows(t2) ; - int backward_cell_number(x, y, z) ; - double backward_weights(t3) ; - int backward_columns(t3) ; - int backward_rows(t2) ; + int backward_cell_number(x, y, z) ; + double backward_weights(t3) ; + int backward_columns(t3) ; + int backward_rows(t2) ; """ + import numpy as np def calc_cell_numbers(maps): - """ Given a field line map dictionary, assign numbers to evolving - cells and boundary cells """ + """Given a field line map dictionary, assign numbers to evolving + cells and boundary cells""" nx, ny, nz = maps["R"].shape MXG = maps["MXG"] # Number of evolving cells @@ -124,7 +125,7 @@ def calc_cell_numbers(maps): def calc_interpolation(cell_number, MXG, yoffset, xtarr, ztarr): """ Calculate CSR format matrix representing a 2D (X-Z) interpolation operation - + Implements the cubic Catmull-Rom spline Coefficients taken from https://en.wikipedia.org/wiki/Cubic_Hermite_spline """ @@ -135,10 +136,10 @@ def calc_interpolation(cell_number, MXG, yoffset, xtarr, ztarr): def weights1D(u): return 0.5 * np.array( [ - -u ** 3 + 2.0 * u ** 2 - u, - 3.0 * u ** 3 - 5.0 * u ** 2 + 2, - -3.0 * u ** 3 + 4.0 * u ** 2 + u, - u ** 3 - u ** 2, + -(u**3) + 2.0 * u**2 - u, + 3.0 * u**3 - 5.0 * u**2 + 2, + -3.0 * u**3 + 4.0 * u**2 + u, + u**3 - u**2, ] ) From 1448576995d2301101a95adf68001eaa73d194d2 Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Sun, 15 Mar 2026 19:59:56 -0700 Subject: [PATCH 03/11] weights: Improving forward/backward weights --- zoidberg/weights.py | 35 ++++++++++++++++++++++++----------- 1 file changed, 24 insertions(+), 11 deletions(-) diff --git a/zoidberg/weights.py b/zoidberg/weights.py index c7c8ac1..3b17e8d 100644 --- a/zoidberg/weights.py +++ b/zoidberg/weights.py @@ -8,10 +8,11 @@ Cell index numbers are stored in three arrays: - cell_number[x,y,z] <- These can be calculated, not stored + cell_number[x,y,z] cell_number_yup[x,y,z] cell_number_ydown[x,y,z] +Non-negative numbers indicate a valid cell number. For forward and backward maps the Nw weights are stored in CSR format @@ -64,7 +65,7 @@ def calc_cell_numbers(maps): N_evolving = (nx - 2 * MXG) * ny * nz # Numbering - cell_number_array = np.zeros((nx, ny, nz), dtype=int) + cell_number_array = np.full((nx, ny, nz), -1, dtype=int) cell_number = 0 for i in range(MXG, nx - MXG): for j in range(ny): @@ -85,8 +86,8 @@ def calc_cell_numbers(maps): cell_number_array[i, j, k] = cell_number cell_number += 1 - backward_cell_number = np.zeros((nx, ny, nz), dtype=int) - forward_cell_number = np.zeros((nx, ny, nz), dtype=int) + backward_cell_number = np.full((nx, ny, nz), -1, dtype=int) + forward_cell_number = np.full((nx, ny, nz), -1, dtype=int) # Iterate through for forward and backward maps forward_xt_prime = maps["forward_xt_prime"] @@ -122,7 +123,7 @@ def calc_cell_numbers(maps): } -def calc_interpolation(cell_number, MXG, yoffset, xtarr, ztarr): +def calc_interpolation(cell_number, MXG, yoffset, xtarr, ztarr, bndry_cell_number): """ Calculate CSR format matrix representing a 2D (X-Z) interpolation operation @@ -144,6 +145,7 @@ def weights1D(u): ) # CSR format + weights = [] columns = [] rows = [] @@ -157,8 +159,11 @@ def weights1D(u): xt = xtarr[i, j, k] zt = ztarr[i, j, k] if (xt < 0.0) or (xt >= nx): - # Boundary - rows.append(-1) + # Boundary => Take value from yup/ydown boundary cell + rows.append(weight_number) + columns.append(bndry_cell_number[i, j, k]) + weights.append(1.0) + weight_number += 1 else: # Not a boundary point => Interpolating rows.append(weight_number) @@ -168,17 +173,23 @@ def weights1D(u): weights_x = weights1D(xt - xi) weights_z = weights1D(zt - zi) + # Weights should sum to 1 + assert abs(np.sum(weights_x) - 1) < 1e-8 + assert abs(np.sum(weights_z) - 1) < 1e-8 + for xo, xw in zip(offsets1D, weights_x): for zo, zw in zip(offsets1D, weights_z): - columns.append( - cell_number[ + col = cell_number[ np.clip(xi + xo, 0, nx - 1), (j + yoffset + ny) % ny, (zi + zo + nz) % nz, - ] - ) + ] + assert col >= 0 + columns.append(col) weights.append(xw * zw) weight_number += 1 + assert weight_number == len(weights) + assert len(weights) == len(columns) return {"weights": weights, "columns": columns, "rows": rows} @@ -197,6 +208,7 @@ def calc_weights(maps): +1, maps["forward_xt_prime"], maps["forward_zt_prime"], + numbering["forward_cell_number"] ) backward = calc_interpolation( @@ -205,6 +217,7 @@ def calc_weights(maps): -1, maps["backward_xt_prime"], maps["backward_zt_prime"], + numbering["backward_cell_number"] ) return { From bffb46b7f1080192a019f61b3f9f7c52f74eb405 Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Mon, 16 Mar 2026 16:53:15 -0700 Subject: [PATCH 04/11] Weights: Add rows in Y boundaries Needed to enable correct calculation of divergence operator. --- zoidberg/weights.py | 182 +++++++++++++++++++++++++++++--------------- 1 file changed, 119 insertions(+), 63 deletions(-) diff --git a/zoidberg/weights.py b/zoidberg/weights.py index 3b17e8d..04eccfe 100644 --- a/zoidberg/weights.py +++ b/zoidberg/weights.py @@ -101,16 +101,16 @@ def calc_cell_numbers(maps): for i in range(MXG, nx - MXG): for j in range(ny): for k in range(nz): - if backward_xt_prime[i, j, k] < 0.0: + if backward_xt_prime[i, j, k] < MXG: backward_cell_number[i, j, k] = cell_number cell_number += 1 - elif backward_xt_prime[i, j, k] >= nx: + elif backward_xt_prime[i, j, k] > nx - MXG - 1: backward_cell_number[i, j, k] = cell_number cell_number += 1 - if forward_xt_prime[i, j, k] < 0.0: + if forward_xt_prime[i, j, k] < MXG: forward_cell_number[i, j, k] = cell_number cell_number += 1 - elif forward_xt_prime[i, j, k] >= nx: + elif forward_xt_prime[i, j, k] > nx - MXG - 1: forward_cell_number[i, j, k] = cell_number cell_number += 1 @@ -123,7 +123,7 @@ def calc_cell_numbers(maps): } -def calc_interpolation(cell_number, MXG, yoffset, xtarr, ztarr, bndry_cell_number): +def calc_interpolation(numbering, maps): """ Calculate CSR format matrix representing a 2D (X-Z) interpolation operation @@ -131,6 +131,9 @@ def calc_interpolation(cell_number, MXG, yoffset, xtarr, ztarr, bndry_cell_numbe Coefficients taken from https://en.wikipedia.org/wiki/Cubic_Hermite_spline """ + cell_number = numbering["cell_number"] + MXG = maps["MXG"] + # Offsets and weights for 1D interpolation offsets1D = [-1, 0, 1, 2] @@ -144,82 +147,135 @@ def weights1D(u): ] ) - # CSR format - - weights = [] - columns = [] - rows = [] + # At the end the operators will be stored in CSR + # format. When assembling the rows are not set in order + # so a more flexible format is needed. + # To use: Append (col, weight) pairs to the row lists. + class Matrix: + def __init__(self, N): + self.rows = [[] for _ in range(N)] + + def csr(self): + """Calculates CSR format arrays""" + rows = np.full(len(self.rows), -1, dtype=int) + weights = [] + columns = [] + for i, entries in enumerate(self.rows): + if len(entries) == 0: + continue + rows[i] = len(weights) + for col, weight in entries: + columns.append(col) + weights.append(weight) + return {"weights": weights, "columns": columns, "rows": rows} + + forward_mat = Matrix(numbering["N_cells"]) + backward_mat = Matrix(numbering["N_cells"]) nx, ny, nz = cell_number.shape - weight_number = 0 # Track location in weights & columns arrays - for i in range(MXG, nx - MXG): - for j in range(ny): - for k in range(nz): - xt = xtarr[i, j, k] - zt = ztarr[i, j, k] - if (xt < 0.0) or (xt >= nx): - # Boundary => Take value from yup/ydown boundary cell - rows.append(weight_number) - columns.append(bndry_cell_number[i, j, k]) - weights.append(1.0) - weight_number += 1 - else: - # Not a boundary point => Interpolating - rows.append(weight_number) - xi = int(xt) # Floor - zi = int(zt) - - weights_x = weights1D(xt - xi) - weights_z = weights1D(zt - zi) - - # Weights should sum to 1 - assert abs(np.sum(weights_x) - 1) < 1e-8 - assert abs(np.sum(weights_z) - 1) < 1e-8 - - for xo, xw in zip(offsets1D, weights_x): - for zo, zw in zip(offsets1D, weights_z): - col = cell_number[ + # Calculate first forward interpolation coefficients then backward + # + # This looks complicated because the rows corresponding to boundary + # cells in the forward direction are set using the backward mapping. + # + # Rows in both forward_mat and backward_mat are set from + # bndry_cell_number + def calc(mat, rev_mat, yoffset, xtarr, ztarr, bndry_cell_number): + for i in range(MXG, nx - MXG): + for j in range(ny): + for k in range(nz): + row = cell_number[i, j, k] + assert row >= 0 + + xt = xtarr[i, j, k] + zt = ztarr[i, j, k] + if (xt < MXG) or (xt > nx - MXG - 1): + # Boundary => Take value from yup/ydown boundary cell + bndry_row = bndry_cell_number[i, j, k] + assert bndry_row >= 0 + mat.rows[row] = [(bndry_row, 1.0)] + + # Boundary cells "free" boundary extrapolate + mat.rows[bndry_row] = [(row, -1.0), (bndry_row, 2.0)] + + # In the reverse matrix e.g. going forward + # from the backward boundary cell. This is + # included because the divergence operator is + # the transpose of the gradient. If the + # divergence in an evolving cell depends on a + # boundary cell, then the gradient in the + # boundary cell must depend on the evolving + # cell. + rev_mat.rows[bndry_row] = [(row, 1.0)] + else: + # Not a boundary point => Interpolating + xi = int(xt) # Floor + zi = int(zt) + + weights_x = weights1D(xt - xi) + weights_z = weights1D(zt - zi) + + # Ensure that points in the X boundary are not used + for wi in range(len(offsets1D)): + if xi + offsets1D[wi] < MXG: + # Move weight away from boundary + weights_x[wi + 1] += weights_x[wi] + weights_x[wi] = 0.0 + for wi in reversed(range(len(offsets1D))): + if xi + offsets1D[wi] >= nx - MXG: + weights_x[wi - 1] += weights_x[wi] + weights_x[wi] = 0.0 + + # Weights should sum to 1 + assert abs(np.sum(weights_x) - 1) < 1e-8 + assert abs(np.sum(weights_z) - 1) < 1e-8 + + for xo, xw in zip(offsets1D, weights_x): + if abs(xw) < 1e-10: + continue + for zo, zw in zip(offsets1D, weights_z): + col = cell_number[ np.clip(xi + xo, 0, nx - 1), (j + yoffset + ny) % ny, (zi + zo + nz) % nz, - ] - assert col >= 0 - columns.append(col) - weights.append(xw * zw) - weight_number += 1 - assert weight_number == len(weights) - assert len(weights) == len(columns) - return {"weights": weights, "columns": columns, "rows": rows} - + ] + assert col >= 0 -def calc_weights(maps): - """ - Calculate interpolation weights for forward and backward maps. + mat.rows[row].append((col, xw * zw)) - Returns a dictionary of arrays to be read into BOUT++ - """ - - numbering = calc_cell_numbers(maps) - - forward = calc_interpolation( - numbering["cell_number"], - maps["MXG"], + calc( + forward_mat, + backward_mat, +1, maps["forward_xt_prime"], maps["forward_zt_prime"], - numbering["forward_cell_number"] + numbering["forward_cell_number"], ) - backward = calc_interpolation( - numbering["cell_number"], - maps["MXG"], + calc( + backward_mat, + forward_mat, -1, maps["backward_xt_prime"], maps["backward_zt_prime"], - numbering["backward_cell_number"] + numbering["backward_cell_number"], ) + return forward_mat.csr(), backward_mat.csr() + + +def calc_weights(maps): + """ + Calculate interpolation weights for forward and backward maps. + + Returns a dictionary of arrays to be read into BOUT++ + """ + + numbering = calc_cell_numbers(maps) + + forward, backward = calc_interpolation(numbering, maps) + return { "cell_number": numbering["cell_number"], "total_cells": numbering["N_cells"], From 889b7a2513a0be13877d3c7b032f619170c62107 Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Mon, 16 Mar 2026 21:22:55 -0700 Subject: [PATCH 05/11] Weights: LLM makeover Renaming, some refactoring, more documentation, standard CSR output. --- zoidberg/weights.py | 710 ++++++++++++++++++++++++++++++-------------- 1 file changed, 483 insertions(+), 227 deletions(-) diff --git a/zoidberg/weights.py b/zoidberg/weights.py index 04eccfe..19d0cf6 100644 --- a/zoidberg/weights.py +++ b/zoidberg/weights.py @@ -1,290 +1,546 @@ -"""Routines to calculate interpolation weights +"""Construct interpolation operators for field-line maps. -N cells divided between Ne evolving cells and Nb = N - Ne boundary cells +This module assigns global cell numbers and builds forward and backward +interpolation operators for BOUT++ mesh files. - [0 .. evolving cells .. (Ne - 1) | Ne .. boundary cells .. (N - 1)] +Numbering scheme +---------------- +Cells are numbered in three arrays: -Boundary cells include both radial (X) boundaries and parallel (Yup/Ydown) cells. +- ``cell_number[x, y, z]``: + Numbering for evolving cells and radial X-boundary cells. +- ``forward_boundary_number[x, y, z]``: + Numbering for forward parallel boundary cells (``yup``). +- ``backward_boundary_number[x, y, z]``: + Numbering for backward parallel boundary cells (``ydown``). -Cell index numbers are stored in three arrays: +A value of ``-1`` means that no numbered cell exists at that location. - cell_number[x,y,z] - cell_number_yup[x,y,z] - cell_number_ydown[x,y,z] +The numbering order is: -Non-negative numbers indicate a valid cell number. +1. evolving cells in the interior ``x = MXG .. nx-MXG-1`` +2. inner radial boundary cells +3. outer radial boundary cells +4. backward parallel boundary cells +5. forward parallel boundary cells -For forward and backward maps the Nw weights are stored in CSR format +Operator storage format +----------------------- +The forward and backward operators are stored as sparse matrices over +all numbered cells and exported in **standard CSR (Compressed Sparse Row)** +format with three arrays: - weights[Nw] - column_index[Nw] - row_index[Ne] <- Starting index into weights and column_index - This will be -1 for boundary points - Needs to be considered when getting the weights +- ``weights[nnz]``: + nonzero matrix entries +- ``columns[nnz]``: + column indices corresponding to ``weights`` +- ``rows[n_rows + 1]``: + CSR row pointer array -For each evolving cell i in 0...(Ne-1) the weight index j is -row_index[i]..(row_index[i+1] - 1) +Row ``i`` occupies entries -i.e. +``weights[rows[i] : rows[i+1]]`` -result[i] = sum_{j = row_index[i]}^{row_index[i+1]-1} weight[j] * input[column_index[j]] +with column indices -The column_index cells go from 0..(N-1), including boundary cells. -Note: row_index[i+1] may be -1, so skip over -1 entries. +``columns[rows[i] : rows[i+1]]``. -Note that these operators are usually represented as non-square -matrices: Input (columns) of length N, output (rows) of length Ne < N. -This is because boundary conditions are set independently. +Empty rows satisfy ``rows[i] == rows[i+1]``. +Interpolation +------------- +Interior interpolation in X-Z uses a 4x4 Catmull-Rom stencil. The Y direction +is shifted by ``+1`` for the forward map and ``-1`` for the backward map. -The output grid file will contain - int cell_number(x, y, z) ; - int total_cells ; +If a mapped point lands in the radial boundary region, the operator uses the +corresponding parallel boundary cell instead of an interior interpolation +stencil. +""" - int forward_cell_number(x, y, z) ; - double forward_weights(t) ; - int forward_columns(t) ; - int forward_rows(t2) ; +import numpy as np - int backward_cell_number(x, y, z) ; - double backward_weights(t3) ; - int backward_columns(t3) ; - int backward_rows(t2) ; -""" +def hits_radial_boundary(x_mapped, nx, mxg): + """Return True if a mapped X coordinate lies in the radial boundary region.""" + return (x_mapped < mxg) or (x_mapped > nx - mxg - 1) -import numpy as np +def assign_cell_numbers(maps): + """Assign global numbers to interior and boundary cells. + + Parameters + ---------- + maps : dict + Field-line map dictionary. Required entries are: + ``R``, ``MXG``, ``forward_xt_prime``, and ``backward_xt_prime``. -def calc_cell_numbers(maps): - """Given a field line map dictionary, assign numbers to evolving - cells and boundary cells""" + Returns + ------- + dict + Dictionary containing: + - ``N_cells``: total number of numbered cells + - ``N_evolving``: number of evolving interior cells + - ``cell_number``: numbering for interior and radial boundary cells + - ``forward_boundary_number``: numbering for forward parallel boundaries + - ``backward_boundary_number``: numbering for backward parallel boundaries + + Notes + ----- + Radial X-boundary cells are numbered directly in ``cell_number``. + Parallel boundary cells are numbered separately in + ``forward_boundary_number`` and ``backward_boundary_number``. + """ nx, ny, nz = maps["R"].shape - MXG = maps["MXG"] - # Number of evolving cells - N_evolving = (nx - 2 * MXG) * ny * nz - - # Numbering - cell_number_array = np.full((nx, ny, nz), -1, dtype=int) - cell_number = 0 - for i in range(MXG, nx - MXG): - for j in range(ny): - for k in range(nz): - cell_number_array[i, j, k] = cell_number - cell_number += 1 + mxg = maps["MXG"] + + n_evolving = (nx - 2 * mxg) * ny * nz + + cell_number = np.full((nx, ny, nz), -1, dtype=int) + + next_cell = 0 + + # Interior evolving cells + for x in range(mxg, nx - mxg): + for y in range(ny): + for z in range(nz): + cell_number[x, y, z] = next_cell + next_cell += 1 # Inner radial boundary cells - for i in range(MXG): - for j in range(ny): - for k in range(nz): - cell_number_array[i, j, k] = cell_number - cell_number += 1 + for x in range(mxg): + for y in range(ny): + for z in range(nz): + cell_number[x, y, z] = next_cell + next_cell += 1 + # Outer radial boundary cells - for i in range(nx - MXG, nx): - for j in range(ny): - for k in range(nz): - cell_number_array[i, j, k] = cell_number - cell_number += 1 + for x in range(nx - mxg, nx): + for y in range(ny): + for z in range(nz): + cell_number[x, y, z] = next_cell + next_cell += 1 - backward_cell_number = np.full((nx, ny, nz), -1, dtype=int) - forward_cell_number = np.full((nx, ny, nz), -1, dtype=int) + backward_boundary_number = np.full((nx, ny, nz), -1, dtype=int) + forward_boundary_number = np.full((nx, ny, nz), -1, dtype=int) - # Iterate through for forward and backward maps forward_xt_prime = maps["forward_xt_prime"] backward_xt_prime = maps["backward_xt_prime"] - # Number of radial boundary cells - N_radial = 2 * MXG * ny * nz - cell_number = N_evolving + N_radial - - # Add the parallel boundary cells - for i in range(MXG, nx - MXG): - for j in range(ny): - for k in range(nz): - if backward_xt_prime[i, j, k] < MXG: - backward_cell_number[i, j, k] = cell_number - cell_number += 1 - elif backward_xt_prime[i, j, k] > nx - MXG - 1: - backward_cell_number[i, j, k] = cell_number - cell_number += 1 - if forward_xt_prime[i, j, k] < MXG: - forward_cell_number[i, j, k] = cell_number - cell_number += 1 - elif forward_xt_prime[i, j, k] > nx - MXG - 1: - forward_cell_number[i, j, k] = cell_number - cell_number += 1 + # Parallel boundary numbering starts after all interior + radial cells + + for x in range(mxg, nx - mxg): + for y in range(ny): + for z in range(nz): + if hits_radial_boundary(backward_xt_prime[x, y, z], nx, mxg): + backward_boundary_number[x, y, z] = next_cell + next_cell += 1 + + if hits_radial_boundary(forward_xt_prime[x, y, z], nx, mxg): + forward_boundary_number[x, y, z] = next_cell + next_cell += 1 return { - "N_cells": cell_number, - "N_evolving": N_evolving, - "cell_number": cell_number_array, - "forward_cell_number": forward_cell_number, - "backward_cell_number": backward_cell_number, + "N_cells": next_cell, + "N_evolving": n_evolving, + "cell_number": cell_number, + "forward_boundary_number": forward_boundary_number, + "backward_boundary_number": backward_boundary_number, } -def calc_interpolation(numbering, maps): +def catmull_rom_weights(u): + """Return 1D Catmull-Rom interpolation weights for fractional offset ``u``. + + Parameters + ---------- + u : float + Fractional coordinate relative to the left stencil point. Typically + in the interval ``[0, 1)``. + + Returns + ------- + numpy.ndarray + Four Catmull-Rom weights corresponding to offsets ``[-1, 0, 1, 2]``. """ - Calculate CSR format matrix representing a 2D (X-Z) interpolation operation + return 0.5 * np.array( + [ + -(u**3) + 2.0 * u**2 - u, + 3.0 * u**3 - 5.0 * u**2 + 2.0, + -3.0 * u**3 + 4.0 * u**2 + u, + u**3 - u**2, + ] + ) + - Implements the cubic Catmull-Rom spline - Coefficients taken from https://en.wikipedia.org/wiki/Cubic_Hermite_spline +def shift_weights_out_of_x_boundaries(weights_x, x_base, x_offsets, nx, mxg): + """Move X interpolation weight off radial boundary points. + + Parameters + ---------- + weights_x : numpy.ndarray + 1D interpolation weights for X offsets. + x_base : int + Floor of the mapped X coordinate. + x_offsets : sequence[int] + Stencil offsets corresponding to ``weights_x``. + nx : int + Number of X points. + mxg : int + Width of the radial boundary region. + + Returns + ------- + numpy.ndarray + Adjusted X weights that do not reference radial boundary cells. + + Notes + ----- + If a stencil point falls into the inner radial boundary, its weight is + added to the nearest point to the right. If it falls into the outer + radial boundary, its weight is added to the nearest point to the left. """ + adjusted = np.array(weights_x, copy=True) - cell_number = numbering["cell_number"] - MXG = maps["MXG"] - - # Offsets and weights for 1D interpolation - offsets1D = [-1, 0, 1, 2] - - def weights1D(u): - return 0.5 * np.array( - [ - -(u**3) + 2.0 * u**2 - u, - 3.0 * u**3 - 5.0 * u**2 + 2, - -3.0 * u**3 + 4.0 * u**2 + u, - u**3 - u**2, - ] + for wi, offset in enumerate(x_offsets): + if x_base + offset < mxg: + adjusted[wi + 1] += adjusted[wi] + adjusted[wi] = 0.0 + + for wi in reversed(range(len(x_offsets))): + offset = x_offsets[wi] + if x_base + offset >= nx - mxg: + adjusted[wi - 1] += adjusted[wi] + adjusted[wi] = 0.0 + + if abs(np.sum(adjusted) - 1.0) >= 1e-8: + raise ValueError(f"Adjusted X weights should sum to 1, got {np.sum(adjusted)}") + + return adjusted + + +class SparseRowMatrixBuilder: + """Accumulate sparse rows before exporting to CSR format. + + Parameters + ---------- + n_rows : int + Number of matrix rows. + + Notes + ----- + Rows are stored internally as Python lists of ``(column, weight)`` pairs. + This allows rows to be assembled out of order during operator construction. + + The exported format is **standard CSR (Compressed Sparse Row)**: + + * ``weights[nnz]`` — nonzero matrix entries + * ``columns[nnz]`` — column indices corresponding to ``weights`` + * ``rows[n_rows + 1]`` — CSR row pointer array + + The CSR invariant is: + + ``rows[i]`` gives the starting index of row ``i`` in ``weights`` and + ``columns``, and ``rows[i+1]`` gives the end index. + + Therefore + + ``rows[i+1] - rows[i]`` + + is the number of nonzero entries in row ``i``. + + Empty rows are represented by + + ``rows[i] == rows[i+1]``. + """ + + def __init__(self, n_rows): + self.rows = [[] for _ in range(n_rows)] + + def set_row(self, row, entries): + """Assign a row exactly once. + + Parameters + ---------- + row : int + Row index to assign. + entries : list[tuple[int, float]] + Sparse row entries. + + Raises + ------ + ValueError + If the row has already been assigned. + """ + if self.rows[row]: + raise ValueError(f"Row {row} assigned more than once") + + self.rows[row] = list(entries) + + def append_to_row(self, row, col, weight): + """Append one entry to a row.""" + self.rows[row].append((col, weight)) + + def to_csr(self): + """Export matrix to standard CSR format using preallocated arrays. + + Returns + ------- + dict + Dictionary containing + + * ``weights`` : numpy.ndarray + Nonzero matrix entries. + * ``columns`` : numpy.ndarray + Column indices corresponding to ``weights``. + * ``rows`` : numpy.ndarray + CSR row pointer array of length ``n_rows + 1``. + + Notes + ----- + This implementation performs two passes: + + 1. Count nonzeros in each row to construct the CSR row pointer. + 2. Allocate arrays of the correct size and fill them. + + This avoids Python list growth during export and is significantly + faster for large matrices. + """ + + n_rows = len(self.rows) + + # ---- pass 1: count nonzeros ---- + + nnz_per_row = np.fromiter( + (len(r) for r in self.rows), + dtype=int, + count=n_rows, ) - # At the end the operators will be stored in CSR - # format. When assembling the rows are not set in order - # so a more flexible format is needed. - # To use: Append (col, weight) pairs to the row lists. - class Matrix: - def __init__(self, N): - self.rows = [[] for _ in range(N)] - - def csr(self): - """Calculates CSR format arrays""" - rows = np.full(len(self.rows), -1, dtype=int) - weights = [] - columns = [] - for i, entries in enumerate(self.rows): - if len(entries) == 0: - continue - rows[i] = len(weights) - for col, weight in entries: - columns.append(col) - weights.append(weight) - return {"weights": weights, "columns": columns, "rows": rows} - - forward_mat = Matrix(numbering["N_cells"]) - backward_mat = Matrix(numbering["N_cells"]) + row_ptr = np.empty(n_rows + 1, dtype=int) + row_ptr[0] = 0 + np.cumsum(nnz_per_row, out=row_ptr[1:]) + + nnz = row_ptr[-1] + + # ---- allocate arrays ---- + + weights = np.empty(nnz, dtype=float) + columns = np.empty(nnz, dtype=int) + + # ---- pass 2: fill arrays ---- + + for row, entries in enumerate(self.rows): + start = row_ptr[row] + + for i, (col, weight) in enumerate(entries): + columns[start + i] = col + weights[start + i] = weight + + return { + "weights": weights, + "columns": columns, + "rows": row_ptr, + } + + +def add_boundary_or_interpolation_row( + matrix_builder, + reverse_builder, + row, + mapped_x, + mapped_z, + y, + y_offset, + boundary_row, + cell_number, + mxg, +): + """Add one operator row for either a boundary map or interior interpolation. + + Parameters + ---------- + matrix_builder : SparseRowMatrixBuilder + Matrix currently being assembled. + reverse_builder : SparseRowMatrixBuilder + Opposite-direction matrix, used to set the reverse boundary row. + row : int + Output row index for the evolving cell. + mapped_x, mapped_z : float + Mapped X and Z coordinates. + y : int + Source Y index. + y_offset : int + Y shift to apply: ``+1`` for forward, ``-1`` for backward. + boundary_row : int + Parallel boundary cell number for this mapped point, or ``-1`` if none. + cell_number : numpy.ndarray + Cell-number array for interior and radial boundary cells. + mxg : int + Radial boundary width. + + Notes + ----- + If the mapped point lands in the radial boundary region, the row is set to + reference the corresponding parallel boundary cell. + + The reverse operator also receives a boundary row contribution. This row is + expected to be unique; if the same reverse boundary row is assigned twice, + an exception is raised. + """ nx, ny, nz = cell_number.shape - # Calculate first forward interpolation coefficients then backward - # - # This looks complicated because the rows corresponding to boundary - # cells in the forward direction are set using the backward mapping. - # - # Rows in both forward_mat and backward_mat are set from - # bndry_cell_number - def calc(mat, rev_mat, yoffset, xtarr, ztarr, bndry_cell_number): - for i in range(MXG, nx - MXG): - for j in range(ny): - for k in range(nz): - row = cell_number[i, j, k] - assert row >= 0 - - xt = xtarr[i, j, k] - zt = ztarr[i, j, k] - if (xt < MXG) or (xt > nx - MXG - 1): - # Boundary => Take value from yup/ydown boundary cell - bndry_row = bndry_cell_number[i, j, k] - assert bndry_row >= 0 - mat.rows[row] = [(bndry_row, 1.0)] - - # Boundary cells "free" boundary extrapolate - mat.rows[bndry_row] = [(row, -1.0), (bndry_row, 2.0)] - - # In the reverse matrix e.g. going forward - # from the backward boundary cell. This is - # included because the divergence operator is - # the transpose of the gradient. If the - # divergence in an evolving cell depends on a - # boundary cell, then the gradient in the - # boundary cell must depend on the evolving - # cell. - rev_mat.rows[bndry_row] = [(row, 1.0)] - else: - # Not a boundary point => Interpolating - xi = int(xt) # Floor - zi = int(zt) - - weights_x = weights1D(xt - xi) - weights_z = weights1D(zt - zi) - - # Ensure that points in the X boundary are not used - for wi in range(len(offsets1D)): - if xi + offsets1D[wi] < MXG: - # Move weight away from boundary - weights_x[wi + 1] += weights_x[wi] - weights_x[wi] = 0.0 - for wi in reversed(range(len(offsets1D))): - if xi + offsets1D[wi] >= nx - MXG: - weights_x[wi - 1] += weights_x[wi] - weights_x[wi] = 0.0 - - # Weights should sum to 1 - assert abs(np.sum(weights_x) - 1) < 1e-8 - assert abs(np.sum(weights_z) - 1) < 1e-8 - - for xo, xw in zip(offsets1D, weights_x): - if abs(xw) < 1e-10: - continue - for zo, zw in zip(offsets1D, weights_z): - col = cell_number[ - np.clip(xi + xo, 0, nx - 1), - (j + yoffset + ny) % ny, - (zi + zo + nz) % nz, - ] - assert col >= 0 - - mat.rows[row].append((col, xw * zw)) - - calc( - forward_mat, - backward_mat, + if hits_radial_boundary(mapped_x, nx, mxg): + if boundary_row < 0: + raise ValueError( + f"Expected a valid boundary row for mapped_x={mapped_x}, got {boundary_row}" + ) + + matrix_builder.set_row(row, [(boundary_row, 1.0)]) + matrix_builder.set_row(boundary_row, [(row, -1.0), (boundary_row, 2.0)]) + reverse_builder.set_row(boundary_row, [(row, 1.0)]) + return + + x_offsets = [-1, 0, 1, 2] + + x_base = int(mapped_x) + z_base = int(mapped_z) + + weights_x = catmull_rom_weights(mapped_x - x_base) + weights_z = catmull_rom_weights(mapped_z - z_base) + + weights_x = shift_weights_out_of_x_boundaries(weights_x, x_base, x_offsets, nx, mxg) + + if abs(np.sum(weights_z) - 1.0) >= 1e-8: + raise ValueError(f"Z weights should sum to 1, got {np.sum(weights_z)}") + + for x_offset, x_weight in zip(x_offsets, weights_x): + if abs(x_weight) < 1e-10: + continue + + for z_offset, z_weight in zip(x_offsets, weights_z): + col = cell_number[ + np.clip(x_base + x_offset, 0, nx - 1), + (y + y_offset + ny) % ny, + (z_base + z_offset + nz) % nz, + ] + + if col < 0: + raise ValueError( + f"Invalid column index at " + f"x={x_base + x_offset}, y={(y + y_offset + ny) % ny}, z={(z_base + z_offset + nz) % nz}" + ) + + matrix_builder.append_to_row(row, col, x_weight * z_weight) + + +def build_parallel_interpolation_operators(numbering, maps): + """Build forward and backward interpolation operators. + + Parameters + ---------- + numbering : dict + Output of :func:`assign_cell_numbers`. + maps : dict + Field-line map dictionary. + + Returns + ------- + tuple[dict, dict] + Two dictionaries in row-start sparse format: + ``(forward_operator, backward_operator)``. + + Notes + ----- + The forward operator shifts in Y by ``+1`` using + ``forward_xt_prime`` / ``forward_zt_prime``. + The backward operator shifts in Y by ``-1`` using + ``backward_xt_prime`` / ``backward_zt_prime``. + """ + cell_number = numbering["cell_number"] + nx, ny, nz = cell_number.shape + mxg = maps["MXG"] + + forward_builder = SparseRowMatrixBuilder(numbering["N_cells"]) + backward_builder = SparseRowMatrixBuilder(numbering["N_cells"]) + + def assemble_direction( + matrix_builder, + reverse_builder, + y_offset, + xt_prime, + zt_prime, + boundary_number, + ): + """Assemble one directional interpolation operator.""" + for x in range(mxg, nx - mxg): + for y in range(ny): + for z in range(nz): + row = cell_number[x, y, z] + if row < 0: + raise ValueError(f"Invalid evolving row at x={x}, y={y}, z={z}") + + add_boundary_or_interpolation_row( + matrix_builder=matrix_builder, + reverse_builder=reverse_builder, + row=row, + mapped_x=xt_prime[x, y, z], + mapped_z=zt_prime[x, y, z], + y=y, + y_offset=y_offset, + boundary_row=boundary_number[x, y, z], + cell_number=cell_number, + mxg=mxg, + ) + + assemble_direction( + forward_builder, + backward_builder, +1, maps["forward_xt_prime"], maps["forward_zt_prime"], - numbering["forward_cell_number"], + numbering["forward_boundary_number"], ) - calc( - backward_mat, - forward_mat, + assemble_direction( + backward_builder, + forward_builder, -1, maps["backward_xt_prime"], maps["backward_zt_prime"], - numbering["backward_cell_number"], + numbering["backward_boundary_number"], ) - return forward_mat.csr(), backward_mat.csr() + return forward_builder.to_csr(), backward_builder.to_csr() -def calc_weights(maps): - """ - Calculate interpolation weights for forward and backward maps. - - Returns a dictionary of arrays to be read into BOUT++ - """ +def calculate_parallel_map_operators(maps): + """Calculate numbering arrays and interpolation operators for BOUT++ + PetscOperators. - numbering = calc_cell_numbers(maps) + Parameters + ---------- + maps : dict + Field-line map dictionary. - forward, backward = calc_interpolation(numbering, maps) + Returns + ------- + dict + Dictionary containing numbering arrays and sparse operator arrays + ready to be written to a BOUT++ mesh file. + """ + numbering = assign_cell_numbers(maps) + forward_op, backward_op = build_parallel_interpolation_operators(numbering, maps) return { "cell_number": numbering["cell_number"], "total_cells": numbering["N_cells"], - "forward_cell_number": numbering["forward_cell_number"], - "forward_weights": forward["weights"], - "forward_columns": forward["columns"], - "forward_rows": forward["rows"], - "backward_cell_number": numbering["backward_cell_number"], - "backward_weights": backward["weights"], - "backward_columns": backward["columns"], - "backward_rows": backward["rows"], + "forward_cell_number": numbering["forward_boundary_number"], + "forward_weights": forward_op["weights"], + "forward_columns": forward_op["columns"], + "forward_rows": forward_op["rows"], + "backward_cell_number": numbering["backward_boundary_number"], + "backward_weights": backward_op["weights"], + "backward_columns": backward_op["columns"], + "backward_rows": backward_op["rows"], } From 51ce9f78eefd2e5c20bc83e686e1d9fd504ef1ef Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Wed, 18 Mar 2026 21:43:03 -0700 Subject: [PATCH 06/11] make_maps: Add cell sub-sampling In addition to following a field-line from cell center, follow a set of points distributed over each cell. Using Gauss-Legendre points to perform averaging over each cell. --- zoidberg/zoidberg.py | 104 ++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 102 insertions(+), 2 deletions(-) diff --git a/zoidberg/zoidberg.py b/zoidberg/zoidberg.py index 3675039..146ee01 100644 --- a/zoidberg/zoidberg.py +++ b/zoidberg/zoidberg.py @@ -45,7 +45,15 @@ def parallel_slice_field_name(field, offset): return "{}_{}{}".format(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, + samples_per_dim: int = 1, + **kwargs, +): """Make the forward and backward FCI maps Parameters @@ -58,6 +66,10 @@ def make_maps(grid, magnetic_field, nslice=1, quiet=False, field_tracer=None, ** Number of parallel slices in each direction quiet : bool Don't display progress bar + samples_per_dim: int, optional + Trace multiple field-lines within each cell. This is the + number of points per dimension so the number of points + per 2D cell is samples_per_dim**2. kwargs Optional arguments for field line tracing, etc. @@ -93,6 +105,30 @@ def make_maps(grid, magnetic_field, nslice=1, quiet=False, field_tracer=None, ** rtol = kwargs.get("rtol", None) + # Sub-sampling within cells + samples_per_dim = int(samples_per_dim) + nsamples = samples_per_dim**2 + if samples_per_dim > 1: + # Calculate quadrature points in 1D + # Using Gauss-Legendre (alpha = beta = 0.0) + # If samples_per_dim is odd then there is a point at cell center + from scipy.special import roots_jacobi + + xs, ws = roots_jacobi(samples_per_dim, 0.0, 0.0) + # xs are in range [-1, 1]; ws are the weights + xs *= 0.5 # Map to [-0.5, 0.5] + ws *= 0.5 # Sum should now be 1 + assert np.isclose(np.sum(ws), 1.0) + + xoffsets, zoffsets = np.meshgrid( + xs, + xs, + indexing="ij", + ) + weights = np.outer(ws, ws).reshape(-1) + xoffsets = xoffsets.reshape(-1) + zoffsets = zoffsets.reshape(-1) + # The field line maps and coordinates, etc. maps = { "R": R, @@ -102,8 +138,11 @@ def make_maps(grid, magnetic_field, nslice=1, quiet=False, field_tracer=None, ** } # A helper data structure that groups the various field line maps along with the offset + # (xt_prime, ztprime)[x,y,z] is the map traced from cell center + # (xt_primes, ztprimes)[x,y,z,s] is a set of s maps for each cell ParallelSlice = namedtuple( - "ParallelSlice", ["offset", "R", "Z", "xt_prime", "zt_prime"] + "ParallelSlice", + ["offset", "R", "Z", "xt_prime", "zt_prime", "xt_primes", "zt_primes"], ) # A list of the above data structures for each offset we want parallel_slices = [] @@ -119,6 +158,13 @@ def make_maps(grid, magnetic_field, nslice=1, quiet=False, field_tracer=None, ** # Initialise the field arrays -- puts them straight into the result dict for field in field_names: maps[field] = np.zeros(shape) + # Sub-sampled maps, multiple samples per cell + for field in [ + parallel_slice_field_name("xt_primes", offset), + parallel_slice_field_name("zt_primes", offset), + ]: + maps[field] = np.zeros(shape + (nsamples,)) + field_names.append(field) # Get the field arrays we just made and wrap them up in our helper tuple fields = map(lambda x: maps[x], field_names) @@ -170,6 +216,60 @@ def make_maps(grid, magnetic_field, nslice=1, quiet=False, field_tracer=None, ** parallel_slice.xt_prime[:, j, :] = xind parallel_slice.zt_prime[:, j, :] = zind + if nsamples == 1: + # Only one sample per cell + parallel_slice.xt_primes[:, j, :, 0] = xind + parallel_slice.zt_primes[:, j, :, 0] = zind + elif pol_slice is None: + # No slice grid, so hit a boundary + # Not calculating intersection R,Z so no need to follow + parallel_slice.xt_primes[:, j, :, :] = -1 + parallel_slice.zt_primes[:, j, :, :] = -1 + else: + # Multiple samples. + # Interpolate R and Z to get starting locations + from scipy import interpolate + + R = pol.R + Z = pol.Z + nx, nz = pol.R.shape + # Wrap around in Z by adding the arrays + Rinterp = interpolate.RegularGridInterpolator( + (np.arange(nx), np.arange(nz + 2) - 1), + np.concatenate((R[:, -1:], R, R[:, 0:1]), axis=-1), + bounds_error=False, + fill_value=None, # Extrapolates + ) + Zinterp = interpolate.RegularGridInterpolator( + (np.arange(nx), np.arange(nz + 2) - 1), + np.concatenate((Z[:, -1:], Z, Z[:, 0:1]), axis=-1), + bounds_error=False, + fill_value=None, # Extrapolates + ) + + # Indices of the cell centers + xinds, zinds = np.meshgrid(np.arange(nx), np.arange(nz), indexing="ij") + + for s, xo, zo, w in zip(range(nsamples), xoffsets, zoffsets, weights): + coord = field_tracer.follow_field_lines( + Rinterp((xinds + xo, zinds)), + Zinterp((xinds + zo, zinds)), + [ycoord, y_slice], + rtol=rtol, + )[1, ...] + xcoord = coord[:, :, 0] + zcoord = coord[:, :, 1] + xind, zind = pol_slice.findIndex(xcoord, zcoord) + + # Check boundary defined by the field + # Note: A fraction of points may hit the boundary + outside = magnetic_field.boundary.outside(xcoord, y_slice, zcoord) + xind[outside] = -1 + zind[outside] = -1 + + parallel_slice.xt_primes[:, j, :, s] = xind + parallel_slice.zt_primes[:, j, :, s] = zind + if (not quiet) and (ny > 1): if tqdm: prog.update() From ded270b1c2f9ddac4acc3c2dc13d071390cc287b Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Thu, 19 Mar 2026 20:40:51 -0700 Subject: [PATCH 07/11] Sub-cell field line tracing and mapping Adds an option to trace multiple field lines per cell (default is 1). Uses Gauss-Legendre points. The resulting maps are used to calculate forward & backward matrices. --- zoidberg/weights.py | 234 ++++++++++++++++++++++++++++++------------- zoidberg/zoidberg.py | 3 + 2 files changed, 167 insertions(+), 70 deletions(-) diff --git a/zoidberg/weights.py b/zoidberg/weights.py index 19d0cf6..b1ba647 100644 --- a/zoidberg/weights.py +++ b/zoidberg/weights.py @@ -60,9 +60,9 @@ import numpy as np -def hits_radial_boundary(x_mapped, nx, mxg): +def any_hits_radial_boundary(x_mapped, nx, mxg): """Return True if a mapped X coordinate lies in the radial boundary region.""" - return (x_mapped < mxg) or (x_mapped > nx - mxg - 1) + return np.any(x_mapped < mxg) or np.any(x_mapped > nx - mxg - 1) def assign_cell_numbers(maps): @@ -123,21 +123,44 @@ def assign_cell_numbers(maps): backward_boundary_number = np.full((nx, ny, nz), -1, dtype=int) forward_boundary_number = np.full((nx, ny, nz), -1, dtype=int) - forward_xt_prime = maps["forward_xt_prime"] - backward_xt_prime = maps["backward_xt_prime"] + # Arrays of mapped indices [x,y,z,s] + # where 's' is the sub-cell node index + forward_xts = maps["forward_xt_primes"] + backward_xts = maps["backward_xt_primes"] - # Parallel boundary numbering starts after all interior + radial cells + # Cell center maps that are used to set boundary conditions + forward_xt = maps["forward_xt_prime"] + backward_xt = maps["backward_xt_prime"] + # Parallel boundary numbering starts after all interior + radial cells for x in range(mxg, nx - mxg): for y in range(ny): for z in range(nz): - if hits_radial_boundary(backward_xt_prime[x, y, z], nx, mxg): + # Allocate a boundary cell if any sub-cell points hit the boundary + if any_hits_radial_boundary(backward_xts[x, y, z, :], nx, mxg): backward_boundary_number[x, y, z] = next_cell next_cell += 1 - - if hits_radial_boundary(forward_xt_prime[x, y, z], nx, mxg): + # Ensure that the cell is marked as a boundary + if not any_hits_radial_boundary(backward_xt[x, y, z], nx, mxg): + # Find a point that hits the boundary + for s in range(backward_xts.shape[-1]): + if any_hits_radial_boundary( + backward_xts[x, y, z, s], nx, mxg + ): + backward_xt[x, y, z] = backward_xts[x, y, z, s] + break + + if any_hits_radial_boundary(forward_xts[x, y, z, :], nx, mxg): forward_boundary_number[x, y, z] = next_cell next_cell += 1 + if not any_hits_radial_boundary(forward_xt[x, y, z], nx, mxg): + # Find a point that hits the boundary + for s in range(forward_xts.shape[-1]): + if any_hits_radial_boundary( + forward_xts[x, y, z, s], nx, mxg + ): + forward_xt[x, y, z] = forward_xts[x, y, z, s] + break return { "N_cells": next_cell, @@ -276,8 +299,14 @@ def set_row(self, row, entries): self.rows[row] = list(entries) - def append_to_row(self, row, col, weight): - """Append one entry to a row.""" + def add_to_entry(self, row, col, weight): + """Add to an entry. If there is already an entry + at this (row, col) then it will be added.""" + for i, (entry_col, entry_weight) in enumerate(self.rows[row]): + if entry_col == col: + # Replace this entry with the sum + self.rows[row][i] = (col, entry_weight + weight) + return self.rows[row].append((col, weight)) def to_csr(self): @@ -344,18 +373,20 @@ def to_csr(self): def add_boundary_or_interpolation_row( - matrix_builder, - reverse_builder, - row, - mapped_x, - mapped_z, - y, - y_offset, - boundary_row, + matrix_builder: SparseRowMatrixBuilder, + reverse_builder: SparseRowMatrixBuilder, + row: int, + mapped_xs, + mapped_zs, + weights, + y: int, + y_offset: int, + boundary_row: int, cell_number, - mxg, + mxg: int, ): - """Add one operator row for either a boundary map or interior interpolation. + """Add one operator row for a boundary map, interior interpolation, + or weighted combination. Parameters ---------- @@ -365,8 +396,10 @@ def add_boundary_or_interpolation_row( Opposite-direction matrix, used to set the reverse boundary row. row : int Output row index for the evolving cell. - mapped_x, mapped_z : float - Mapped X and Z coordinates. + mapped_xs, mapped_xs : [float] + Mapped X and Z coordinates for all sub-cell node indices. + weights : [float] + Weighting factor for each sub-cell node. Sums to 1. y : int Source Y index. y_offset : int @@ -378,60 +411,96 @@ def add_boundary_or_interpolation_row( mxg : int Radial boundary width. + Returns + ------- + boundary_weight: float + A factor between 0 and 1 that should be applied to the volume + of boundary cells when constructing divergence operators. + It is the fraction of the cell that hits the boundary. + Notes ----- - If the mapped point lands in the radial boundary region, the row is set to - reference the corresponding parallel boundary cell. + If one or more mapped points land in the radial boundary region, + the row is set to reference the corresponding parallel boundary cell. The reverse operator also receives a boundary row contribution. This row is expected to be unique; if the same reverse boundary row is assigned twice, an exception is raised. """ + assert len(weights) == len(mapped_xs) + assert len(weights) == len(mapped_zs) + if abs(np.sum(weights) - 1.0) >= 1e-8: + raise ValueError(f"weights should sum to 1, got {np.sum(weights)}") nx, ny, nz = cell_number.shape - if hits_radial_boundary(mapped_x, nx, mxg): - if boundary_row < 0: - raise ValueError( - f"Expected a valid boundary row for mapped_x={mapped_x}, got {boundary_row}" - ) + x_offsets = [-1, 0, 1, 2] - matrix_builder.set_row(row, [(boundary_row, 1.0)]) - matrix_builder.set_row(boundary_row, [(row, -1.0), (boundary_row, 2.0)]) - reverse_builder.set_row(boundary_row, [(row, 1.0)]) - return + # Track whether any sub-cell nodes hit boundaries + any_hit_boundary = False + # Sum of the weights of points hitting the boundary + boundary_weight = 0.0 - x_offsets = [-1, 0, 1, 2] + # Iterate over sub-cell nodes + for mapped_x, mapped_z, subcell_weight in zip(mapped_xs, mapped_zs, weights): + if any_hits_radial_boundary(mapped_x, nx, mxg): + any_hit_boundary = True + boundary_weight += subcell_weight + continue + x_base = int(mapped_x) + z_base = int(mapped_z) - x_base = int(mapped_x) - z_base = int(mapped_z) + weights_x = catmull_rom_weights(mapped_x - x_base) + weights_z = catmull_rom_weights(mapped_z - z_base) - weights_x = catmull_rom_weights(mapped_x - x_base) - weights_z = catmull_rom_weights(mapped_z - z_base) + weights_x = shift_weights_out_of_x_boundaries( + weights_x, x_base, x_offsets, nx, mxg + ) - weights_x = shift_weights_out_of_x_boundaries(weights_x, x_base, x_offsets, nx, mxg) + if abs(np.sum(weights_z) - 1.0) >= 1e-8: + raise ValueError(f"Z weights should sum to 1, got {np.sum(weights_z)}") - if abs(np.sum(weights_z) - 1.0) >= 1e-8: - raise ValueError(f"Z weights should sum to 1, got {np.sum(weights_z)}") + for x_offset, x_weight in zip(x_offsets, weights_x): + if abs(x_weight) < 1e-10: + continue - for x_offset, x_weight in zip(x_offsets, weights_x): - if abs(x_weight) < 1e-10: - continue + for z_offset, z_weight in zip(x_offsets, weights_z): + col = cell_number[ + np.clip(x_base + x_offset, 0, nx - 1), + (y + y_offset + ny) % ny, + (z_base + z_offset + nz) % nz, + ] + + if col < 0: + raise ValueError( + f"Invalid column index at " + f"x={x_base + x_offset}, y={(y + y_offset + ny) % ny}, z={(z_base + z_offset + nz) % nz}" + ) - for z_offset, z_weight in zip(x_offsets, weights_z): - col = cell_number[ - np.clip(x_base + x_offset, 0, nx - 1), - (y + y_offset + ny) % ny, - (z_base + z_offset + nz) % nz, - ] - - if col < 0: - raise ValueError( - f"Invalid column index at " - f"x={x_base + x_offset}, y={(y + y_offset + ny) % ny}, z={(z_base + z_offset + nz) % nz}" + # Sub-cell nodes will provide entries at overlapping + # (row, col) so add to existing entries. + matrix_builder.add_to_entry( + row, col, subcell_weight * x_weight * z_weight ) - matrix_builder.append_to_row(row, col, x_weight * z_weight) + if any_hit_boundary: + # One or more points hit a boundary + if boundary_row < 0: + raise ValueError( + f"Expected a valid boundary row for mapped_x={mapped_xs}, got {boundary_row}" + ) + + matrix_builder.add_to_entry(row, boundary_row, boundary_weight) + # Note: Identity operation in the boundary + # Extrapolating messes up the boundary fluxes! + matrix_builder.set_row( + boundary_row, + [ + (boundary_row, 1.0), + ], + ) + reverse_builder.set_row(boundary_row, [(row, 1.0)]) + return boundary_weight def build_parallel_interpolation_operators(numbering, maps): @@ -447,8 +516,11 @@ def build_parallel_interpolation_operators(numbering, maps): Returns ------- tuple[dict, dict] - Two dictionaries in row-start sparse format: - ``(forward_operator, backward_operator)``. + Two dictionaries in row-start sparse format, + and 3D arrays of volume fractions for forward + and backward boundaries. + ``(forward_operator, backward_operator, + forward_fraction, backward_fraction)``. Notes ----- @@ -464,13 +536,18 @@ def build_parallel_interpolation_operators(numbering, maps): forward_builder = SparseRowMatrixBuilder(numbering["N_cells"]) backward_builder = SparseRowMatrixBuilder(numbering["N_cells"]) + forward_boundary_fraction = np.zeros(cell_number.shape) + backward_boundary_fraction = np.zeros(cell_number.shape) + def assemble_direction( matrix_builder, reverse_builder, y_offset, - xt_prime, - zt_prime, + xt_primes, + zt_primes, + weights, boundary_number, + boundary_fraction, ): """Assemble one directional interpolation operator.""" for x in range(mxg, nx - mxg): @@ -480,12 +557,14 @@ def assemble_direction( if row < 0: raise ValueError(f"Invalid evolving row at x={x}, y={y}, z={z}") - add_boundary_or_interpolation_row( + # Store the fraction [0,1] of the cell that hits + boundary_fraction[x, y, z] = add_boundary_or_interpolation_row( matrix_builder=matrix_builder, reverse_builder=reverse_builder, row=row, - mapped_x=xt_prime[x, y, z], - mapped_z=zt_prime[x, y, z], + mapped_xs=xt_primes[x, y, z, :], + mapped_zs=zt_primes[x, y, z, :], + weights=weights, y=y, y_offset=y_offset, boundary_row=boundary_number[x, y, z], @@ -497,21 +576,30 @@ def assemble_direction( forward_builder, backward_builder, +1, - maps["forward_xt_prime"], - maps["forward_zt_prime"], + maps["forward_xt_primes"], + maps["forward_zt_primes"], + maps["subcell_weights"], numbering["forward_boundary_number"], + forward_boundary_fraction, ) assemble_direction( backward_builder, forward_builder, -1, - maps["backward_xt_prime"], - maps["backward_zt_prime"], + maps["backward_xt_primes"], + maps["backward_zt_primes"], + maps["subcell_weights"], numbering["backward_boundary_number"], + backward_boundary_fraction, ) - return forward_builder.to_csr(), backward_builder.to_csr() + return ( + forward_builder.to_csr(), + backward_builder.to_csr(), + forward_boundary_fraction, + backward_boundary_fraction, + ) def calculate_parallel_map_operators(maps): @@ -521,7 +609,9 @@ def calculate_parallel_map_operators(maps): Parameters ---------- maps : dict - Field-line map dictionary. + Field-line map dictionary. Uses `forward_{x,z}t_primes` and + `backward_{x,z}t_primes` if present, falling back to + `forward_{x,z}t_prime` and `backward_{x,z}t_prime`. Returns ------- @@ -530,16 +620,20 @@ def calculate_parallel_map_operators(maps): ready to be written to a BOUT++ mesh file. """ numbering = assign_cell_numbers(maps) - forward_op, backward_op = build_parallel_interpolation_operators(numbering, maps) + forward_op, backward_op, forward_fraction, backward_fraction = ( + build_parallel_interpolation_operators(numbering, maps) + ) return { "cell_number": numbering["cell_number"], "total_cells": numbering["N_cells"], "forward_cell_number": numbering["forward_boundary_number"], + "forward_boundary_fraction": forward_fraction, "forward_weights": forward_op["weights"], "forward_columns": forward_op["columns"], "forward_rows": forward_op["rows"], "backward_cell_number": numbering["backward_boundary_number"], + "backward_boundary_fraction": backward_fraction, "backward_weights": backward_op["weights"], "backward_columns": backward_op["columns"], "backward_rows": backward_op["rows"], diff --git a/zoidberg/zoidberg.py b/zoidberg/zoidberg.py index 146ee01..f2875d6 100644 --- a/zoidberg/zoidberg.py +++ b/zoidberg/zoidberg.py @@ -128,6 +128,8 @@ def make_maps( weights = np.outer(ws, ws).reshape(-1) xoffsets = xoffsets.reshape(-1) zoffsets = zoffsets.reshape(-1) + else: + weights = np.ones(1) # The field line maps and coordinates, etc. maps = { @@ -135,6 +137,7 @@ def make_maps( "Z": Z, "MXG": mxg, "MYG": nslice, + "subcell_weights": weights, } # A helper data structure that groups the various field line maps along with the offset From 3193ea06998e42815c98959a83d5a86ab9e04172 Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Sun, 22 Mar 2026 17:04:39 -0700 Subject: [PATCH 08/11] Weights: Major rewrite, calculating C, L+ and L- numberings Cell center values, forward legs and backward legs are now three separate 'spaces' with different numberings. The forward and backward operators are now non-square matrices that map between spaces. This complexity is required to handle cells that partly intersect boundaries: Fluxes are split between the interior and boundary part. Such cells therefore have two forward (or backward) 'legs' on which fluxes are calculated. --- zoidberg/weights.py | 821 +++++++++++++++++++++----------------------- 1 file changed, 390 insertions(+), 431 deletions(-) diff --git a/zoidberg/weights.py b/zoidberg/weights.py index b1ba647..df6d197 100644 --- a/zoidberg/weights.py +++ b/zoidberg/weights.py @@ -1,166 +1,232 @@ -"""Construct interpolation operators for field-line maps. +"""Construct cell-space numbering, leg metadata, and leg operators. -This module assigns global cell numbers and builds forward and backward -interpolation operators for BOUT++ mesh files. +This module assigns global numbers for the cell space ``C`` and constructs +metadata and operators for directional leg spaces: -Numbering scheme ----------------- -Cells are numbered in three arrays: +- ``L+``: forward legs +- ``L-``: backward legs -- ``cell_number[x, y, z]``: - Numbering for evolving cells and radial X-boundary cells. -- ``forward_boundary_number[x, y, z]``: - Numbering for forward parallel boundary cells (``yup``). -- ``backward_boundary_number[x, y, z]``: - Numbering for backward parallel boundary cells (``ydown``). +The cell space ``C`` contains: -A value of ``-1`` means that no numbered cell exists at that location. - -The numbering order is: - -1. evolving cells in the interior ``x = MXG .. nx-MXG-1`` -2. inner radial boundary cells -3. outer radial boundary cells +1. evolving interior cells +2. inner radial X-boundary cells +3. outer radial X-boundary cells 4. backward parallel boundary cells 5. forward parallel boundary cells -Operator storage format ------------------------ -The forward and backward operators are stored as sparse matrices over -all numbered cells and exported in **standard CSR (Compressed Sparse Row)** -format with three arrays: - -- ``weights[nnz]``: - nonzero matrix entries -- ``columns[nnz]``: - column indices corresponding to ``weights`` -- ``rows[n_rows + 1]``: - CSR row pointer array +The forward and backward leg spaces contain rows only for evolving interior +source cells. Each interior cell contributes either one or two rows in each leg +space: -Row ``i`` occupies entries +- one row if all sub-samples map to interior cells, or all map to the boundary +- two rows if some sub-samples map to interior cells and others hit the + boundary -``weights[rows[i] : rows[i+1]]`` +For each direction this module exports: -with column indices +- explicit leg-number metadata arrays +- a leg-weight vector containing the quadrature weight of each leg row +- a destination operator mapping from cell space ``C`` to leg space + (rows = leg rows, columns = cell indices) -``columns[rows[i] : rows[i+1]]``. +No injection matrices are generated. The consumer can reconstruct local +cell-to-leg injection operators directly from the leg-number metadata. -Empty rows satisfy ``rows[i] == rows[i+1]``. +Sparse storage format +--------------------- +Sparse matrices are exported in standard CSR (Compressed Sparse Row) format: -Interpolation -------------- -Interior interpolation in X-Z uses a 4x4 Catmull-Rom stencil. The Y direction -is shifted by ``+1`` for the forward map and ``-1`` for the backward map. +- ``weights[nnz]``: nonzero matrix entries +- ``columns[nnz]``: column indices corresponding to ``weights`` +- ``rows[n_rows + 1]``: CSR row pointer array -If a mapped point lands in the radial boundary region, the operator uses the -corresponding parallel boundary cell instead of an interior interpolation -stencil. +Row ``i`` occupies entries ``weights[rows[i]:rows[i+1]]`` and +``columns[rows[i]:rows[i+1]]``. """ +from __future__ import annotations + import numpy as np +FORWARD_DIRECTION = { + "name": "forward", + "xt_key": "forward_xt_primes", + "xt_fallback": "forward_xt_prime", + "zt_key": "forward_zt_primes", + "zt_fallback": "forward_zt_prime", + "boundary_number_key": "forward_boundary_number", + "y_offset": +1, +} + +BACKWARD_DIRECTION = { + "name": "backward", + "xt_key": "backward_xt_primes", + "xt_fallback": "backward_xt_prime", + "zt_key": "backward_zt_primes", + "zt_fallback": "backward_zt_prime", + "boundary_number_key": "backward_boundary_number", + "y_offset": -1, +} + + +def _get_map_array(maps, primary_key, fallback_key=None): + """Return an array from the map dictionary.""" + if primary_key in maps: + return maps[primary_key] + if fallback_key is not None and fallback_key in maps: + return maps[fallback_key] + if fallback_key is None: + raise KeyError(f"Missing required map array '{primary_key}'") + raise KeyError( + f"Missing required map array '{primary_key}' " + f"(fallback '{fallback_key}' also not found)" + ) + -def any_hits_radial_boundary(x_mapped, nx, mxg): +def _as_subsample_array(array): + """Ensure mapped coordinate arrays have shape ``(nx, ny, nz, n_subsamples)``.""" + arr = np.asarray(array) + if arr.ndim == 3: + return arr[..., np.newaxis] + if arr.ndim == 4: + return arr + raise ValueError( + f"Expected mapped coordinate array with 3 or 4 dimensions, got shape {arr.shape}" + ) + + +def _get_subcell_weights(maps, n_subsamples): + """Return the quadrature weights for sub-samples as a 1D array.""" + if "subcell_weights" in maps: + weights = np.asarray(maps["subcell_weights"], dtype=float) + elif n_subsamples == 1: + weights = np.full(1, 1.0, dtype=float) + else: + raise ValueError("Missing subcell_weights") + + if weights.ndim != 1: + raise ValueError( + f"subcell_weights must be one-dimensional, got shape {weights.shape}" + ) + if len(weights) != n_subsamples: + raise ValueError( + f"subcell_weights has length {len(weights)} but expected {n_subsamples}" + ) + if abs(np.sum(weights) - 1.0) >= 1e-8: + raise ValueError(f"subcell_weights should sum to 1, got {np.sum(weights)}") + return weights + + +def hits_radial_boundary(x_mapped, nx, mxg): """Return True if a mapped X coordinate lies in the radial boundary region.""" - return np.any(x_mapped < mxg) or np.any(x_mapped > nx - mxg - 1) + return (x_mapped < mxg) or (x_mapped > nx - mxg - 1) -def assign_cell_numbers(maps): - """Assign global numbers to interior and boundary cells. +def classify_leg_contributions(mapped_xs, subcell_weights, nx, mxg): + """Classify one source cell's sub-samples into interior and boundary parts. + + Parameters + ---------- + mapped_xs : array-like + Forward or backward mapped X coordinates for one source cell. + subcell_weights : array-like + Quadrature weights for the sub-samples. + nx : int + Number of X points. + mxg : int + Width of the radial boundary region. + + Returns + ------- + dict + Classification information for one source cell. + """ + mapped_xs = np.asarray(mapped_xs, dtype=float) + subcell_weights = np.asarray(subcell_weights, dtype=float) + + hits_boundary_mask = np.fromiter( + (hits_radial_boundary(x, nx, mxg) for x in mapped_xs), + dtype=bool, + count=len(mapped_xs), + ) + + boundary_fraction = float(np.sum(subcell_weights[hits_boundary_mask])) + interior_fraction = float(np.sum(subcell_weights[~hits_boundary_mask])) + + return { + "hits_boundary": hits_boundary_mask, + "has_boundary_leg": np.any(hits_boundary_mask), + "has_interior_leg": np.any(~hits_boundary_mask), + "boundary_fraction": boundary_fraction, + "interior_fraction": interior_fraction, + } + + +def assign_cell_space_numbers(maps): + """Assign global numbers to cell-space interior and boundary cells. Parameters ---------- maps : dict - Field-line map dictionary. Required entries are: - ``R``, ``MXG``, ``forward_xt_prime``, and ``backward_xt_prime``. + Field-line map dictionary. Returns ------- dict - Dictionary containing: - - ``N_cells``: total number of numbered cells - - ``N_evolving``: number of evolving interior cells - - ``cell_number``: numbering for interior and radial boundary cells - - ``forward_boundary_number``: numbering for forward parallel boundaries - - ``backward_boundary_number``: numbering for backward parallel boundaries - - Notes - ----- - Radial X-boundary cells are numbered directly in ``cell_number``. - Parallel boundary cells are numbered separately in - ``forward_boundary_number`` and ``backward_boundary_number``. + Cell-space numbering metadata. """ nx, ny, nz = maps["R"].shape mxg = maps["MXG"] n_evolving = (nx - 2 * mxg) * ny * nz - cell_number = np.full((nx, ny, nz), -1, dtype=int) next_cell = 0 - # Interior evolving cells + # Interior cells for x in range(mxg, nx - mxg): for y in range(ny): for z in range(nz): cell_number[x, y, z] = next_cell next_cell += 1 - # Inner radial boundary cells + # Inner X boundary for x in range(mxg): for y in range(ny): for z in range(nz): cell_number[x, y, z] = next_cell next_cell += 1 - # Outer radial boundary cells + # Outer X boundary for x in range(nx - mxg, nx): for y in range(ny): for z in range(nz): cell_number[x, y, z] = next_cell next_cell += 1 + forward_xts = _as_subsample_array( + _get_map_array(maps, "forward_xt_primes", "forward_xt_prime") + ) + backward_xts = _as_subsample_array( + _get_map_array(maps, "backward_xt_primes", "backward_xt_prime") + ) + backward_boundary_number = np.full((nx, ny, nz), -1, dtype=int) forward_boundary_number = np.full((nx, ny, nz), -1, dtype=int) - # Arrays of mapped indices [x,y,z,s] - # where 's' is the sub-cell node index - forward_xts = maps["forward_xt_primes"] - backward_xts = maps["backward_xt_primes"] - - # Cell center maps that are used to set boundary conditions - forward_xt = maps["forward_xt_prime"] - backward_xt = maps["backward_xt_prime"] - - # Parallel boundary numbering starts after all interior + radial cells for x in range(mxg, nx - mxg): for y in range(ny): for z in range(nz): - # Allocate a boundary cell if any sub-cell points hit the boundary - if any_hits_radial_boundary(backward_xts[x, y, z, :], nx, mxg): + if np.any( + [hits_radial_boundary(v, nx, mxg) for v in backward_xts[x, y, z, :]] + ): backward_boundary_number[x, y, z] = next_cell next_cell += 1 - # Ensure that the cell is marked as a boundary - if not any_hits_radial_boundary(backward_xt[x, y, z], nx, mxg): - # Find a point that hits the boundary - for s in range(backward_xts.shape[-1]): - if any_hits_radial_boundary( - backward_xts[x, y, z, s], nx, mxg - ): - backward_xt[x, y, z] = backward_xts[x, y, z, s] - break - - if any_hits_radial_boundary(forward_xts[x, y, z, :], nx, mxg): + if np.any( + [hits_radial_boundary(v, nx, mxg) for v in forward_xts[x, y, z, :]] + ): forward_boundary_number[x, y, z] = next_cell next_cell += 1 - if not any_hits_radial_boundary(forward_xt[x, y, z], nx, mxg): - # Find a point that hits the boundary - for s in range(forward_xts.shape[-1]): - if any_hits_radial_boundary( - forward_xts[x, y, z, s], nx, mxg - ): - forward_xt[x, y, z] = forward_xts[x, y, z, s] - break return { "N_cells": next_cell, @@ -172,19 +238,7 @@ def assign_cell_numbers(maps): def catmull_rom_weights(u): - """Return 1D Catmull-Rom interpolation weights for fractional offset ``u``. - - Parameters - ---------- - u : float - Fractional coordinate relative to the left stencil point. Typically - in the interval ``[0, 1)``. - - Returns - ------- - numpy.ndarray - Four Catmull-Rom weights corresponding to offsets ``[-1, 0, 1, 2]``. - """ + """Return 1D Catmull-Rom interpolation weights for fractional offset ``u``.""" return 0.5 * np.array( [ -(u**3) + 2.0 * u**2 - u, @@ -196,32 +250,7 @@ def catmull_rom_weights(u): def shift_weights_out_of_x_boundaries(weights_x, x_base, x_offsets, nx, mxg): - """Move X interpolation weight off radial boundary points. - - Parameters - ---------- - weights_x : numpy.ndarray - 1D interpolation weights for X offsets. - x_base : int - Floor of the mapped X coordinate. - x_offsets : sequence[int] - Stencil offsets corresponding to ``weights_x``. - nx : int - Number of X points. - mxg : int - Width of the radial boundary region. - - Returns - ------- - numpy.ndarray - Adjusted X weights that do not reference radial boundary cells. - - Notes - ----- - If a stencil point falls into the inner radial boundary, its weight is - added to the nearest point to the right. If it falls into the outer - radial boundary, its weight is added to the nearest point to the left. - """ + """Move X interpolation weight off radial boundary points.""" adjusted = np.array(weights_x, copy=True) for wi, offset in enumerate(x_offsets): @@ -242,217 +271,184 @@ def shift_weights_out_of_x_boundaries(weights_x, x_base, x_offsets, nx, mxg): class SparseRowMatrixBuilder: - """Accumulate sparse rows before exporting to CSR format. - - Parameters - ---------- - n_rows : int - Number of matrix rows. - - Notes - ----- - Rows are stored internally as Python lists of ``(column, weight)`` pairs. - This allows rows to be assembled out of order during operator construction. - - The exported format is **standard CSR (Compressed Sparse Row)**: - - * ``weights[nnz]`` — nonzero matrix entries - * ``columns[nnz]`` — column indices corresponding to ``weights`` - * ``rows[n_rows + 1]`` — CSR row pointer array - - The CSR invariant is: - - ``rows[i]`` gives the starting index of row ``i`` in ``weights`` and - ``columns``, and ``rows[i+1]`` gives the end index. - - Therefore - - ``rows[i+1] - rows[i]`` - - is the number of nonzero entries in row ``i``. - - Empty rows are represented by - - ``rows[i] == rows[i+1]``. - """ + """Accumulate sparse rows before exporting to standard CSR format.""" def __init__(self, n_rows): self.rows = [[] for _ in range(n_rows)] def set_row(self, row, entries): - """Assign a row exactly once. - - Parameters - ---------- - row : int - Row index to assign. - entries : list[tuple[int, float]] - Sparse row entries. - - Raises - ------ - ValueError - If the row has already been assigned. - """ + """Assign a row exactly once.""" if self.rows[row]: raise ValueError(f"Row {row} assigned more than once") - self.rows[row] = list(entries) def add_to_entry(self, row, col, weight): - """Add to an entry. If there is already an entry - at this (row, col) then it will be added.""" + """Add to a matrix entry, accumulating duplicate columns in a row.""" for i, (entry_col, entry_weight) in enumerate(self.rows[row]): if entry_col == col: - # Replace this entry with the sum self.rows[row][i] = (col, entry_weight + weight) return self.rows[row].append((col, weight)) def to_csr(self): - """Export matrix to standard CSR format using preallocated arrays. - - Returns - ------- - dict - Dictionary containing - - * ``weights`` : numpy.ndarray - Nonzero matrix entries. - * ``columns`` : numpy.ndarray - Column indices corresponding to ``weights``. - * ``rows`` : numpy.ndarray - CSR row pointer array of length ``n_rows + 1``. - - Notes - ----- - This implementation performs two passes: - - 1. Count nonzeros in each row to construct the CSR row pointer. - 2. Allocate arrays of the correct size and fill them. - - This avoids Python list growth during export and is significantly - faster for large matrices. - """ - + """Export matrix to standard CSR format using preallocated arrays.""" n_rows = len(self.rows) - - # ---- pass 1: count nonzeros ---- - - nnz_per_row = np.fromiter( - (len(r) for r in self.rows), - dtype=int, - count=n_rows, - ) - + nnz_per_row = np.fromiter((len(r) for r in self.rows), dtype=int, count=n_rows) row_ptr = np.empty(n_rows + 1, dtype=int) row_ptr[0] = 0 np.cumsum(nnz_per_row, out=row_ptr[1:]) nnz = row_ptr[-1] - - # ---- allocate arrays ---- - weights = np.empty(nnz, dtype=float) columns = np.empty(nnz, dtype=int) - # ---- pass 2: fill arrays ---- - for row, entries in enumerate(self.rows): start = row_ptr[row] - for i, (col, weight) in enumerate(entries): columns[start + i] = col weights[start + i] = weight - return { - "weights": weights, - "columns": columns, - "rows": row_ptr, - } + return {"weights": weights, "columns": columns, "rows": row_ptr} -def add_boundary_or_interpolation_row( - matrix_builder: SparseRowMatrixBuilder, - reverse_builder: SparseRowMatrixBuilder, - row: int, - mapped_xs, - mapped_zs, - weights, - y: int, - y_offset: int, - boundary_row: int, - cell_number, - mxg: int, -): - """Add one operator row for a boundary map, interior interpolation, - or weighted combination. +def build_leg_space_metadata(cell_space, maps, direction): + """Build leg numbering metadata for one direction. Parameters ---------- - matrix_builder : SparseRowMatrixBuilder - Matrix currently being assembled. - reverse_builder : SparseRowMatrixBuilder - Opposite-direction matrix, used to set the reverse boundary row. - row : int - Output row index for the evolving cell. - mapped_xs, mapped_xs : [float] - Mapped X and Z coordinates for all sub-cell node indices. - weights : [float] - Weighting factor for each sub-cell node. Sums to 1. - y : int - Source Y index. - y_offset : int - Y shift to apply: ``+1`` for forward, ``-1`` for backward. - boundary_row : int - Parallel boundary cell number for this mapped point, or ``-1`` if none. - cell_number : numpy.ndarray - Cell-number array for interior and radial boundary cells. - mxg : int - Radial boundary width. + cell_space : dict + Output of :func:`assign_cell_space_numbers`. + maps : dict + Field-line map dictionary. + direction : dict + Direction configuration, e.g. ``FORWARD_DIRECTION``. Returns ------- - boundary_weight: float - A factor between 0 and 1 that should be applied to the volume - of boundary cells when constructing divergence operators. - It is the fraction of the cell that hits the boundary. - - Notes - ----- - If one or more mapped points land in the radial boundary region, - the row is set to reference the corresponding parallel boundary cell. - - The reverse operator also receives a boundary row contribution. This row is - expected to be unique; if the same reverse boundary row is assigned twice, - an exception is raised. + dict + Leg-space numbering metadata for the chosen direction. """ - assert len(weights) == len(mapped_xs) - assert len(weights) == len(mapped_zs) - if abs(np.sum(weights) - 1.0) >= 1e-8: - raise ValueError(f"weights should sum to 1, got {np.sum(weights)}") - + cell_number = cell_space["cell_number"] nx, ny, nz = cell_number.shape + mxg = maps["MXG"] + + xt_primes = _as_subsample_array( + _get_map_array(maps, direction["xt_key"], direction["xt_fallback"]) + ) + n_subsamples = xt_primes.shape[3] + subcell_weights = _get_subcell_weights(maps, n_subsamples) + + interior_leg_number = np.full((nx, ny, nz), -1, dtype=int) + boundary_leg_number = np.full((nx, ny, nz), -1, dtype=int) + interior_fraction = np.zeros((nx, ny, nz), dtype=float) + boundary_fraction = np.zeros((nx, ny, nz), dtype=float) + has_interior_leg = np.zeros((nx, ny, nz), dtype=bool) + has_boundary_leg = np.zeros((nx, ny, nz), dtype=bool) + + next_leg = 0 + for x in range(mxg, nx - mxg): + for y in range(ny): + for z in range(nz): + classification = classify_leg_contributions( + xt_primes[x, y, z, :], subcell_weights, nx, mxg + ) + + has_interior_leg[x, y, z] = classification["has_interior_leg"] + has_boundary_leg[x, y, z] = classification["has_boundary_leg"] + interior_fraction[x, y, z] = classification["interior_fraction"] + boundary_fraction[x, y, z] = classification["boundary_fraction"] + + if classification["has_interior_leg"]: + interior_leg_number[x, y, z] = next_leg + next_leg += 1 + if classification["has_boundary_leg"]: + boundary_leg_number[x, y, z] = next_leg + next_leg += 1 + + return { + "direction_name": direction["name"], + "n_legs": next_leg, + "interior_leg_number": interior_leg_number, + "boundary_leg_number": boundary_leg_number, + "interior_fraction": interior_fraction, + "boundary_fraction": boundary_fraction, + "has_interior_leg": has_interior_leg, + "has_boundary_leg": has_boundary_leg, + } + + +def build_leg_weight_vector(leg_space): + """Build the leg-weight vector for one leg space. + + Parameters + ---------- + leg_space : dict + Output of :func:`build_leg_space_metadata`. + + Returns + ------- + numpy.ndarray + One weight per leg row. + """ + weights = np.zeros(leg_space["n_legs"], dtype=float) + + interior_leg_number = leg_space["interior_leg_number"] + boundary_leg_number = leg_space["boundary_leg_number"] + interior_fraction = leg_space["interior_fraction"] + boundary_fraction = leg_space["boundary_fraction"] + has_interior_leg = leg_space["has_interior_leg"] + has_boundary_leg = leg_space["has_boundary_leg"] + + nx, ny, nz = interior_leg_number.shape + for x in range(nx): + for y in range(ny): + for z in range(nz): + if not (has_interior_leg[x, y, z] or has_boundary_leg[x, y, z]): + continue + + split = has_interior_leg[x, y, z] and has_boundary_leg[x, y, z] + + if has_interior_leg[x, y, z]: + row = interior_leg_number[x, y, z] + weights[row] = interior_fraction[x, y, z] if split else 1.0 + if has_boundary_leg[x, y, z]: + row = boundary_leg_number[x, y, z] + weights[row] = boundary_fraction[x, y, z] if split else 1.0 + + return weights + +def accumulate_interior_stencil( + matrix_builder, + row, + mapped_xs, + mapped_zs, + subcell_weights, + hits_boundary_mask, + y, + y_offset, + cell_number, + mxg, +): + """Accumulate the normalized interior stencil for one leg row.""" + nx, ny, nz = cell_number.shape x_offsets = [-1, 0, 1, 2] - # Track whether any sub-cell nodes hit boundaries - any_hit_boundary = False - # Sum of the weights of points hitting the boundary - boundary_weight = 0.0 + interior_weight = float(np.sum(subcell_weights[~hits_boundary_mask])) + if interior_weight <= 0.0: + raise ValueError("Cannot build interior leg row with zero interior weight") - # Iterate over sub-cell nodes - for mapped_x, mapped_z, subcell_weight in zip(mapped_xs, mapped_zs, weights): - if any_hits_radial_boundary(mapped_x, nx, mxg): - any_hit_boundary = True - boundary_weight += subcell_weight + for mapped_x, mapped_z, subcell_weight, hit_boundary in zip( + mapped_xs, mapped_zs, subcell_weights, hits_boundary_mask + ): + if hit_boundary: continue + x_base = int(mapped_x) z_base = int(mapped_z) weights_x = catmull_rom_weights(mapped_x - x_base) weights_z = catmull_rom_weights(mapped_z - z_base) - weights_x = shift_weights_out_of_x_boundaries( weights_x, x_base, x_offsets, nx, mxg ) @@ -460,181 +456,144 @@ def add_boundary_or_interpolation_row( if abs(np.sum(weights_z) - 1.0) >= 1e-8: raise ValueError(f"Z weights should sum to 1, got {np.sum(weights_z)}") + scale = subcell_weight / interior_weight + for x_offset, x_weight in zip(x_offsets, weights_x): if abs(x_weight) < 1e-10: continue - for z_offset, z_weight in zip(x_offsets, weights_z): col = cell_number[ np.clip(x_base + x_offset, 0, nx - 1), (y + y_offset + ny) % ny, (z_base + z_offset + nz) % nz, ] - if col < 0: raise ValueError( - f"Invalid column index at " - f"x={x_base + x_offset}, y={(y + y_offset + ny) % ny}, z={(z_base + z_offset + nz) % nz}" + "Invalid column index at " + f"x={x_base + x_offset}, y={(y + y_offset + ny) % ny}, " + f"z={(z_base + z_offset + nz) % nz}" ) - - # Sub-cell nodes will provide entries at overlapping - # (row, col) so add to existing entries. - matrix_builder.add_to_entry( - row, col, subcell_weight * x_weight * z_weight - ) - - if any_hit_boundary: - # One or more points hit a boundary - if boundary_row < 0: - raise ValueError( - f"Expected a valid boundary row for mapped_x={mapped_xs}, got {boundary_row}" - ) - - matrix_builder.add_to_entry(row, boundary_row, boundary_weight) - # Note: Identity operation in the boundary - # Extrapolating messes up the boundary fluxes! - matrix_builder.set_row( - boundary_row, - [ - (boundary_row, 1.0), - ], - ) - reverse_builder.set_row(boundary_row, [(row, 1.0)]) - return boundary_weight + matrix_builder.add_to_entry(row, col, scale * x_weight * z_weight) -def build_parallel_interpolation_operators(numbering, maps): - """Build forward and backward interpolation operators. +def build_leg_destination_operator(cell_space, leg_space, maps, direction): + """Build the destination operator for one direction. - Parameters - ---------- - numbering : dict - Output of :func:`assign_cell_numbers`. - maps : dict - Field-line map dictionary. - - Returns - ------- - tuple[dict, dict] - Two dictionaries in row-start sparse format, - and 3D arrays of volume fractions for forward - and backward boundaries. - ``(forward_operator, backward_operator, - forward_fraction, backward_fraction)``. - - Notes - ----- - The forward operator shifts in Y by ``+1`` using - ``forward_xt_prime`` / ``forward_zt_prime``. - The backward operator shifts in Y by ``-1`` using - ``backward_xt_prime`` / ``backward_zt_prime``. + The operator has rows in the chosen leg space and columns in cell space. + Boundary leg rows map to one boundary cell in ``C``. Interior leg rows map + to a normalized interpolation stencil over interior/X-boundary cells in + ``C``. """ - cell_number = numbering["cell_number"] + cell_number = cell_space["cell_number"] + boundary_number = cell_space[direction["boundary_number_key"]] nx, ny, nz = cell_number.shape mxg = maps["MXG"] - forward_builder = SparseRowMatrixBuilder(numbering["N_cells"]) - backward_builder = SparseRowMatrixBuilder(numbering["N_cells"]) - - forward_boundary_fraction = np.zeros(cell_number.shape) - backward_boundary_fraction = np.zeros(cell_number.shape) - - def assemble_direction( - matrix_builder, - reverse_builder, - y_offset, - xt_primes, - zt_primes, - weights, - boundary_number, - boundary_fraction, - ): - """Assemble one directional interpolation operator.""" - for x in range(mxg, nx - mxg): - for y in range(ny): - for z in range(nz): - row = cell_number[x, y, z] - if row < 0: - raise ValueError(f"Invalid evolving row at x={x}, y={y}, z={z}") - - # Store the fraction [0,1] of the cell that hits - boundary_fraction[x, y, z] = add_boundary_or_interpolation_row( - matrix_builder=matrix_builder, - reverse_builder=reverse_builder, + xt_primes = _as_subsample_array( + _get_map_array(maps, direction["xt_key"], direction["xt_fallback"]) + ) + zt_primes = _as_subsample_array( + _get_map_array(maps, direction["zt_key"], direction["zt_fallback"]) + ) + subcell_weights = _get_subcell_weights(maps, xt_primes.shape[3]) + + builder = SparseRowMatrixBuilder(leg_space["n_legs"]) + + for x in range(mxg, nx - mxg): + for y in range(ny): + for z in range(nz): + mapped_xs = xt_primes[x, y, z, :] + mapped_zs = zt_primes[x, y, z, :] + classification = classify_leg_contributions( + mapped_xs, subcell_weights, nx, mxg + ) + + if leg_space["has_boundary_leg"][x, y, z]: + row = leg_space["boundary_leg_number"][x, y, z] + col = boundary_number[x, y, z] + if col < 0: + raise ValueError( + f"Expected valid {direction['name']} boundary number at x={x}, y={y}, z={z}" + ) + builder.set_row(row, [(col, 1.0)]) + + if leg_space["has_interior_leg"][x, y, z]: + row = leg_space["interior_leg_number"][x, y, z] + accumulate_interior_stencil( + matrix_builder=builder, row=row, - mapped_xs=xt_primes[x, y, z, :], - mapped_zs=zt_primes[x, y, z, :], - weights=weights, + mapped_xs=mapped_xs, + mapped_zs=mapped_zs, + subcell_weights=subcell_weights, + hits_boundary_mask=classification["hits_boundary"], y=y, - y_offset=y_offset, - boundary_row=boundary_number[x, y, z], + y_offset=direction["y_offset"], cell_number=cell_number, mxg=mxg, ) - assemble_direction( - forward_builder, - backward_builder, - +1, - maps["forward_xt_primes"], - maps["forward_zt_primes"], - maps["subcell_weights"], - numbering["forward_boundary_number"], - forward_boundary_fraction, - ) - - assemble_direction( - backward_builder, - forward_builder, - -1, - maps["backward_xt_primes"], - maps["backward_zt_primes"], - maps["subcell_weights"], - numbering["backward_boundary_number"], - backward_boundary_fraction, - ) - - return ( - forward_builder.to_csr(), - backward_builder.to_csr(), - forward_boundary_fraction, - backward_boundary_fraction, - ) + return builder.to_csr() def calculate_parallel_map_operators(maps): - """Calculate numbering arrays and interpolation operators for BOUT++ - PetscOperators. + """Calculate cell numbering, leg metadata, leg weights, + and forward/backward operators. Parameters ---------- maps : dict - Field-line map dictionary. Uses `forward_{x,z}t_primes` and - `backward_{x,z}t_primes` if present, falling back to - `forward_{x,z}t_prime` and `backward_{x,z}t_prime`. + Field-line map dictionary. Returns ------- dict - Dictionary containing numbering arrays and sparse operator arrays - ready to be written to a BOUT++ mesh file. + Dictionary containing cell-space numbering, forward/backward leg + numbering metadata, leg-weight vectors, and the CSR arrays for + ``forward`` and ``backward``. """ - numbering = assign_cell_numbers(maps) - forward_op, backward_op, forward_fraction, backward_fraction = ( - build_parallel_interpolation_operators(numbering, maps) + cell_space = assign_cell_space_numbers(maps) + + forward_leg_space = build_leg_space_metadata(cell_space, maps, FORWARD_DIRECTION) + backward_leg_space = build_leg_space_metadata(cell_space, maps, BACKWARD_DIRECTION) + + forward_operator = build_leg_destination_operator( + cell_space, forward_leg_space, maps, FORWARD_DIRECTION + ) + backward_operator = build_leg_destination_operator( + cell_space, backward_leg_space, maps, BACKWARD_DIRECTION ) return { - "cell_number": numbering["cell_number"], - "total_cells": numbering["N_cells"], - "forward_cell_number": numbering["forward_boundary_number"], - "forward_boundary_fraction": forward_fraction, - "forward_weights": forward_op["weights"], - "forward_columns": forward_op["columns"], - "forward_rows": forward_op["rows"], - "backward_cell_number": numbering["backward_boundary_number"], - "backward_boundary_fraction": backward_fraction, - "backward_weights": backward_op["weights"], - "backward_columns": backward_op["columns"], - "backward_rows": backward_op["rows"], + # Cell numbering + "total_cells": cell_space["N_cells"], + "cell_number": cell_space["cell_number"], + "forward_cell_number": cell_space["forward_boundary_number"], + "backward_cell_number": cell_space["backward_boundary_number"], + # Forward leg numbering + "n_forward_legs": forward_leg_space["n_legs"], + "forward_leg_interior_number": forward_leg_space["interior_leg_number"], + "forward_leg_boundary_number": forward_leg_space["boundary_leg_number"], + "forward_leg_weights": build_leg_weight_vector(forward_leg_space), + # Forward interpolation operator in CSR format + "forward_weights": forward_operator["weights"], + "forward_columns": forward_operator["columns"], + "forward_rows": forward_operator["rows"], + # Backward leg numbering + "n_backward_legs": backward_leg_space["n_legs"], + "backward_leg_interior_number": backward_leg_space["interior_leg_number"], + "backward_leg_boundary_number": backward_leg_space["boundary_leg_number"], + "backward_leg_weights": build_leg_weight_vector(backward_leg_space), + # Backward interpolation operator in CSR format + "backward_weights": backward_operator["weights"], + "backward_columns": backward_operator["columns"], + "backward_rows": backward_operator["rows"], } + + +__all__ = [ + "assign_cell_space_numbers", + "build_leg_space_metadata", + "build_leg_weight_vector", + "build_leg_destination_operator", + "calculate_parallel_map_operators", +] From 5760dedafff4e8595be72e540ccb09579c0f4709 Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Wed, 13 May 2026 15:27:25 -0700 Subject: [PATCH 09/11] Weights: Ensure consistency with boundaries When cells partly intersect boundaries, parallel boundary conditions need to be applied to the partial boundaries. --- examples/rotating-ellipse.py | 90 ++++++++++++++++++++++++++++++++++++ zoidberg/weights.py | 78 +++++++++++++++++++++++++++++++ 2 files changed, 168 insertions(+) create mode 100755 examples/rotating-ellipse.py diff --git a/examples/rotating-ellipse.py b/examples/rotating-ellipse.py new file mode 100755 index 0000000..ee3a8b4 --- /dev/null +++ b/examples/rotating-ellipse.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python + +# Create grids for a rotating ellipse stellarator, based on curvilinear grids + +import numpy as np + +import zoidberg + +############################################################################# +# Define the magnetic field + +# Length in y after which the coils return to their starting (R,Z) locations +yperiod = 2.0 * np.pi / 5 + +xcentre = 2.0 + +magnetic_field = zoidberg.field.RotatingEllipse( + xcentre=xcentre, # Major radius of axis [m] + radius=0.8, # Minor radius of the coils [m] + I_coil=0.4, # Coil current + yperiod=yperiod, + Btor=1.0, # Toroidal magnetic field +) + +############################################################################# +# Create the inner flux surface, starting at a point at phi=0 +# To do this we need to define the y locations of the poloidal points +# where we will construct grids + +start_r = xcentre + 0.4 +start_z = 0.0 + +nslices = 8 # Number of poloidal slices +ycoords = np.linspace(0, yperiod, nslices) +npoints = 20 # Points per poloidal slice + +rzcoord, ycoords = zoidberg.fieldtracer.trace_poincare( + magnetic_field, start_r, start_z, yperiod, y_slices=ycoords, revs=npoints, nover=1 +) + +inner_lines = [] +for i in range(nslices): + r = rzcoord[:, i, 0, 0] + z = rzcoord[:, i, 0, 1] + line = zoidberg.rzline.line_from_points(r, z) + # Re-map the points so they're approximately uniform in distance along the surface + # Note that this results in some motion of the line + line = line.equallySpaced() + inner_lines.append(line) + +# Now have a list of y coordinates (ycoords) and inner lines (inner_lines) + +############################################################################# +# Generate a fixed circle for the outer boundary + +outer_line = zoidberg.rzline.circle(R0=xcentre, r=0.6) + +############################################################################# +# 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) + for inner_line in inner_lines +] + +############################################################################# +# Create a grid, then calculate forward and backward maps + +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=samples_per_dim) + +maps = zoidberg.weights.modify_maps(maps, partial_boundaries = 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}'") +zoidberg.write_maps( + grid, magnetic_field, maps, gridfile=filename, new_names=False, metric2d=False +) diff --git a/zoidberg/weights.py b/zoidberg/weights.py index df6d197..38509e3 100644 --- a/zoidberg/weights.py +++ b/zoidberg/weights.py @@ -410,9 +410,13 @@ def build_leg_weight_vector(leg_space): if has_interior_leg[x, y, z]: row = interior_leg_number[x, y, z] + if row < 0: + raise IndexError(f"interior_leg_number[{x}, {y}, {z}] < 0") weights[row] = interior_fraction[x, y, z] if split else 1.0 if has_boundary_leg[x, y, z]: row = boundary_leg_number[x, y, z] + if row < 0: + raise IndexError(f"boundary_leg_number[{x}, {y}, {z}] < 0") weights[row] = boundary_fraction[x, y, z] if split else 1.0 return weights @@ -589,6 +593,79 @@ def calculate_parallel_map_operators(maps): "backward_rows": backward_operator["rows"], } +def modify_maps(maps, partial_boundaries = True): + """ + Modify the forward/backward maps, adding forward/backward weights. + + Ensures consistency of boundaries and forward/backward operators, + depending on whether cells can partially intersect boundaries. + + Parameters + ---------- + maps : dict + Field-line map dictionary. + + """ + mxg = maps["MXG"] + if partial_boundaries: + # If any part of the map hits a boundary then it must be marked + # as a boundary cell so that applyBoundary() will work correctly. + for dirname in ['forward', 'backward']: + xts = maps[dirname + '_xt_primes'] + xt = maps[dirname + '_xt_prime'] + + nx, ny, nz, nw = xts.shape + for i in range(nx): + for j in range(ny): + for k in range(nz): + for w in range(nw): + if hits_radial_boundary(xts[i,j,k,w], nx, mxg): + xt[i,j,k] = -1 + break + maps[dirname + '_xt_prime'] = xt + + else: + # Filter maps to remove partial boundary intersection + for dirname in ['forward', 'backward']: + xts = maps[dirname + '_xt_primes'] + xt = maps[dirname + '_xt_prime'] + zts = maps[dirname + '_zt_primes'] + zt = maps[dirname + '_zt_prime'] + + nx, ny, nz, nw = xts.shape + for i in range(nx): + for j in range(ny): + for k in range(nz): + hit_bndry = [hits_radial_boundary(xts[i,j,k,w], nx, mxg) for w in range(nw)] + if np.any(hit_bndry) and not np.all(hit_bndry): + # Partial intersection + # Set boundary if center point hits boundary + if hits_radial_boundary(xt[i,j,k], nx, mxg): + # Center point hits boundary + # Set all points to the same + xts[i,j,k,:] = xt[i,j,k] + zts[i,j,k,:] = zt[i,j,k] + else: + # Center brand does not hit boundary + # Set boundary points to the center value + xts[i,j,k,hit_bndry] = xt[i,j,k] + zts[i,j,k,hit_bndry] = zt[i,j,k] + maps[dirname + '_xt_primes'] = xts + maps[dirname + '_zt_primes'] = zts + + weight_vars = calculate_parallel_map_operators(maps) + maps.update(weight_vars) + # Remove sub-sampled maps from output + for var in [ + "forward_xt_primes", + "forward_zt_primes", + "backward_xt_primes", + "backward_zt_primes", + "subcell_weights", + ]: + maps.pop(var, None) + + return maps __all__ = [ "assign_cell_space_numbers", @@ -596,4 +673,5 @@ def calculate_parallel_map_operators(maps): "build_leg_weight_vector", "build_leg_destination_operator", "calculate_parallel_map_operators", + "modify_maps", ] From bd2a1230791c6c1b30266b95261af7edf6eca318 Mon Sep 17 00:00:00 2001 From: bendudson Date: Wed, 13 May 2026 22:31:02 +0000 Subject: [PATCH 10/11] Apply black/isort changes --- zoidberg/weights.py | 55 ++++++++++++++++++++++++--------------------- 1 file changed, 30 insertions(+), 25 deletions(-) diff --git a/zoidberg/weights.py b/zoidberg/weights.py index 38509e3..031a773 100644 --- a/zoidberg/weights.py +++ b/zoidberg/weights.py @@ -593,7 +593,8 @@ def calculate_parallel_map_operators(maps): "backward_rows": backward_operator["rows"], } -def modify_maps(maps, partial_boundaries = True): + +def modify_maps(maps, partial_boundaries=True): """ Modify the forward/backward maps, adding forward/backward weights. @@ -610,63 +611,67 @@ def modify_maps(maps, partial_boundaries = True): if partial_boundaries: # If any part of the map hits a boundary then it must be marked # as a boundary cell so that applyBoundary() will work correctly. - for dirname in ['forward', 'backward']: - xts = maps[dirname + '_xt_primes'] - xt = maps[dirname + '_xt_prime'] + for dirname in ["forward", "backward"]: + xts = maps[dirname + "_xt_primes"] + xt = maps[dirname + "_xt_prime"] nx, ny, nz, nw = xts.shape for i in range(nx): for j in range(ny): for k in range(nz): for w in range(nw): - if hits_radial_boundary(xts[i,j,k,w], nx, mxg): - xt[i,j,k] = -1 + if hits_radial_boundary(xts[i, j, k, w], nx, mxg): + xt[i, j, k] = -1 break - maps[dirname + '_xt_prime'] = xt + maps[dirname + "_xt_prime"] = xt else: # Filter maps to remove partial boundary intersection - for dirname in ['forward', 'backward']: - xts = maps[dirname + '_xt_primes'] - xt = maps[dirname + '_xt_prime'] - zts = maps[dirname + '_zt_primes'] - zt = maps[dirname + '_zt_prime'] + for dirname in ["forward", "backward"]: + xts = maps[dirname + "_xt_primes"] + xt = maps[dirname + "_xt_prime"] + zts = maps[dirname + "_zt_primes"] + zt = maps[dirname + "_zt_prime"] nx, ny, nz, nw = xts.shape for i in range(nx): for j in range(ny): for k in range(nz): - hit_bndry = [hits_radial_boundary(xts[i,j,k,w], nx, mxg) for w in range(nw)] + hit_bndry = [ + hits_radial_boundary(xts[i, j, k, w], nx, mxg) + for w in range(nw) + ] if np.any(hit_bndry) and not np.all(hit_bndry): # Partial intersection # Set boundary if center point hits boundary - if hits_radial_boundary(xt[i,j,k], nx, mxg): + if hits_radial_boundary(xt[i, j, k], nx, mxg): # Center point hits boundary # Set all points to the same - xts[i,j,k,:] = xt[i,j,k] - zts[i,j,k,:] = zt[i,j,k] + xts[i, j, k, :] = xt[i, j, k] + zts[i, j, k, :] = zt[i, j, k] else: # Center brand does not hit boundary # Set boundary points to the center value - xts[i,j,k,hit_bndry] = xt[i,j,k] - zts[i,j,k,hit_bndry] = zt[i,j,k] - maps[dirname + '_xt_primes'] = xts - maps[dirname + '_zt_primes'] = zts + xts[i, j, k, hit_bndry] = xt[i, j, k] + zts[i, j, k, hit_bndry] = zt[i, j, k] + maps[dirname + "_xt_primes"] = xts + maps[dirname + "_zt_primes"] = zts weight_vars = calculate_parallel_map_operators(maps) maps.update(weight_vars) # Remove sub-sampled maps from output for var in [ - "forward_xt_primes", - "forward_zt_primes", - "backward_xt_primes", - "backward_zt_primes", - "subcell_weights", + "forward_xt_primes", + "forward_zt_primes", + "backward_xt_primes", + "backward_zt_primes", + "subcell_weights", ]: maps.pop(var, None) return maps + __all__ = [ "assign_cell_space_numbers", "build_leg_space_metadata", From 1f694488cac22cc2835eb08513f5c50f1c24db4b Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Tue, 19 May 2026 15:13:58 -0700 Subject: [PATCH 11/11] Ruff formatting --- examples/rotating-ellipse.py | 2 +- zoidberg/zoidberg.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/rotating-ellipse.py b/examples/rotating-ellipse.py index ee3a8b4..6bf8e06 100755 --- a/examples/rotating-ellipse.py +++ b/examples/rotating-ellipse.py @@ -77,7 +77,7 @@ 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=partial_boundaries) ############################################################################# # Write grid file diff --git a/zoidberg/zoidberg.py b/zoidberg/zoidberg.py index 3f86a34..9ca5b81 100644 --- a/zoidberg/zoidberg.py +++ b/zoidberg/zoidberg.py @@ -184,8 +184,8 @@ def make_maps( maps[field] = np.zeros(shape) # Sub-sampled maps, multiple samples per cell for field in [ - parallel_slice_field_name("xt_primes", offset), - parallel_slice_field_name("zt_primes", offset), + parallel_slice_field_name("xt_primes", offset), + parallel_slice_field_name("zt_primes", offset), ]: maps[field] = np.zeros(shape + (nsamples,)) field_names.append(field)