Skip to content

Add selectors and boundary flow source diagnostics #27

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
mikekryjak wants to merge 4 commits into from
Closed

Add selectors and boundary flow source diagnostics #27

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
53c5e0e
Add toroidal cell length, rename dl to dpol
mikekryjak Feb 6, 2026
a0ec8e7
Add midplane selection for double null only
mikekryjak Feb 6, 2026
f49e4ec
Add selectors and guard checks
mikekryjak Feb 6, 2026
f03e374
Add boundary flow source calculation
mikekryjak Feb 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
101 changes: 92 additions & 9 deletions xhermes/accessors.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -220,6 +220,8 @@ def extract_2d_tokamak_geometry(self):
# Cell areas in real space - comes from Jacobian
# Note: can be calculated from flux space or real space coordinates:
# dV = (hthe/Bpol) * (R*Bpol*dr) * dy*2pi = hthe * dy * dr * 2pi * R

# Radial cell length
ds["dr"] = (["x", "theta"], ds.dx.data / (ds.R.data * ds.Bpxy.data))
ds["dr"].attrs.update({
"conversion" : 1,
Expand All @@ -228,24 +230,105 @@ def extract_2d_tokamak_geometry(self):
"long_name" : "Length of cell in the radial direction",
"source" : "xHermes"})

ds["hthe"] = (["x", "theta"], ds["J"].data * ds["Bpxy"].data) # h_theta
ds["hthe"] = (["x", "theta"], ds["J"].data * ds["Bpxy"].data)
ds["hthe"].attrs.update({
"conversion" : 1,
"units" : "m/radian",
"standard_name" : "h_theta: poloidal arc length per radian",
"long_name" : "h_theta: poloidal arc length per radian",
"source" : "xHermes"})

ds["dl"] = (["x", "theta"], ds["dy"].data * ds["hthe"].data) # poloidal arc length
ds["dl"].attrs.update({
"conversion" : 1,
"units" : "m",
"standard_name" : "poloidal arc length",
"long_name" : "Poloidal arc length",
"source" : "xHermes"})


# Toroidal cell length
ds["dtor"] = (
["x", "theta"],
ds["dz"].data * np.sqrt(ds["g_33"].data),
)
ds["dtor"].attrs.update(
{
"conversion": 1,
"units": "m",
"standard_name": "Toroidal length",
"long_name": "Toroidal length",
"source": "xHermes",
}
)

# Poloidal cell length
ds["dpol"] = (
["x", "theta"],
ds["dy"].data * ds["hthe"].data,
) # Poloidal length
ds["dpol"].attrs.update(
{
"conversion": 1,
"units": "m",
"standard_name": "Poloidal length",
"long_name": "Poloidal length",
"source": "xHermes",
}
)

return ds

def calculate_boundary_flows(self):
"""
Calculate particle sources at SOL and PFR boundaries

For each species in the dataset, calculates the particle source
at the SOL and PFR boundaries by dividing the flow diagnostics
by cell volume. These are then added back to the dataset.
This allows boundary flows to be treated like regular sources
which helps make post-processing more straightforward.

TODO: Add energy flows
TODO: Add flows in y direction

Returns
-------
xarray.Dataset
Dataset with added boundary flow source terms
"""

ds = self.data

sol = ds.hermes.select_region("sol_boundary")
pfr = ds.hermes.select_region("pfr_boundary")
sol_guard = ds.hermes.select_region("sol_boundary_guard")

for name in ds.metadata["species"]:

# Account for inconsistency in diagnostic names
if name in ds.metadata["neutral_species"]:
flow_diagnostic_name = f"pf{name}_adv_perp_xlow"
else:
flow_diagnostic_name = f"pf{name}_tot_xlow"

if flow_diagnostic_name in ds.data_vars:

# xlow means SOL boundary flow is read at guard cell
ds[f"S{name}_sol_boundary"] = (
sol_guard[flow_diagnostic_name] / sol["dv"]
)
ds[f"S{name}_sol_boundary"].attrs.update(
{
"short_name": "Particle source",
"long_name": f"Particle source of {name} at SOL boundary",
"units": "m^-3 s^-1",
}
)

# xlow means PFR boundary flow is read at final domain cell
ds[f"S{name}_pfr_boundary"] = pfr[flow_diagnostic_name] / pfr["dv"]
ds[f"S{name}_pfr_boundary"].attrs.update(
{
"short_name": "Particle source",
"long_name": f"Particle source of {name} at PFR boundary",
"units": "m^-3 s^-1",
}
)

return ds

@register_dataarray_accessor("hermes")
class HermesDataArrayAccessor(BoutDataArrayAccessor):
Expand Down
68 changes: 51 additions & 17 deletions xhermes/selectors.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,13 +47,20 @@ def slice_poloidal(ds, name):
else:
index["inner_target"] = nyg - MYG - 1
index["outer_target"] = MYG

elif "double-null" in topology:
index["inner_lower_target"] = MYG
index["outer_lower_target"] = nyg - MYG - 1
index["inner_upper_target"] = ny_innerg - MYG * 2 -1
index["outer_upper_target"] = ny_innerg

# Midplane is inbetween two cell centres, "a" and "b"
# "a" is the lower midplane index, "b" is the upper midplane index
index["outer_midplane_a"] = int((j2_2g - j1_2g) / 2) + j1_2g
index["outer_midplane_b"] = int((j2_2g - j1_2g) / 2) + j1_2g + 1
index["inner_midplane_a"] = int((j2_1g - j1_1g) / 2) + j1_1g
index["inner_midplane_b"] = int((j2_1g - j1_1g) / 2) + j1_1g + 1

# Guard selection
if "guard" in name:
if MYG > 0:
Expand Down Expand Up @@ -112,10 +119,13 @@ def slice_2d(ds, name):
topology = m["topology"]

polidx = lambda x: slice_poloidal(ds, x)
slice_x_domain = slice(MXG, nxg - MXG) # Domain X points (excl guards)
slice_x_outer = nxg - MXG - 1 # Last domain cell on SOL edge side
slice_x_inner = MXG


# Radial selection slices
slice_x_domain = slice(MXG, nxg - MXG)
slice_x_outer_final = nxg - MXG - 1
slice_x_inner_final = MXG
slice_x_outer_guard = nxg - MXG
slice_x_inner_guard = MXG - 1 if MXG == 2 else MXG

slices = {}

Expand All @@ -131,17 +141,17 @@ def slice_2d(ds, name):
np.r_[polidx("inner_target"),
polidx("outer_target")])

slices["core_boundary"] = (slice_x_inner,
slices["core_boundary"] = (slice_x_inner_final,
slice(j1_1g+1, j2_2g+1))

## Lower single null
if "lower" in topology:
slices["sol_boundary"] = (
slice_x_outer,
slice_x_outer_final,
slice(polidx("inner_target"), polidx("outer_target"))
)
slices["pfr_boundary"] = (
slice_x_inner,
slice_x_inner_final,
np.r_[
slice(MYG, j1_1g+1),
slice(j2_2g+1, nyg - MYG)
Expand All @@ -151,11 +161,11 @@ def slice_2d(ds, name):
## Upper single null
elif "upper" in topology:
slices["sol_boundary"] = (
slice_x_outer,
slice_x_outer_final,
slice(polidx("outer_target"), polidx("inner_target"))
)
slices["pfr_boundary"] = (
slice_x_inner,
slice_x_inner_final,
np.r_[
slice(MYG, j1_1g+1),
slice(j2_2g+1, nyg - MYG)
Expand All @@ -176,31 +186,47 @@ def slice_2d(ds, name):
])

slices["sol_inner_boundary"] = (
slice_x_outer,
slice_x_outer_final,
np.r_[
slice(polidx("inner_lower_target"), polidx("inner_upper_target")+1),
])

slices["sol_outer_boundary"] = (
slice_x_outer,
slice_x_outer_final,
np.r_[
slice(polidx("outer_upper_target"), polidx("outer_lower_target")+1),
])

slices["sol_boundary"] = (
slice_x_outer,
slice_x_outer_final,
np.r_[
slice(polidx("inner_lower_target"), polidx("inner_upper_target")+1),
slice(polidx("outer_upper_target"), polidx("outer_lower_target")+1),
])

slices["core_boundary"] = (slice_x_inner,
slices["sol_boundary_guard"] = (
slice_x_outer_guard,
np.r_[
slice(polidx("inner_lower_target"), polidx("inner_upper_target")+1),
slice(polidx("outer_upper_target"), polidx("outer_lower_target")+1),
])


slices["core_boundary"] = (slice_x_inner_final,
np.r_[
slice(j1_1g+1, j2_1g+1),
slice(j1_2g+1, j2_2g+1)]
)

slices["pfr_boundary"] = (slice_x_inner,
slices["pfr_boundary"] = (slice_x_inner_final,
np.r_[
slice(MYG, j1_1g+1),
slice(j2_1g+1, ny_innerg-MYG*2),
slice(ny_innerg, j1_2g+1),
slice(j2_2g+1, nyg-MYG)]
)

slices["pfr_boundary_guard"] = (slice_x_inner_guard,
np.r_[
slice(MYG, j1_1g+1),
slice(j2_1g+1, ny_innerg-MYG*2),
Expand All @@ -210,13 +236,21 @@ def slice_2d(ds, name):

else:
raise ValueError(f"Unknown topology: {topology}")

# Poloidal selections
if name in ["inner_target", "outer_target",
"inner_lower_target", "inner_upper_target",
"outer_lower_target", "outer_upper_target",
"yguards"]:
"yguards",
"inner_midplane_a", "inner_midplane_b",
"outer_midplane_a", "outer_midplane_b"]:

slices[name] = (slice_x_domain, slice_poloidal(ds, name))

if name in ["yguards"] and MYG == 0:
raise ValueError("No y guards found in dataset, cannot select poloidal guards!")

if name in ["sol_boundary_guard", "pfr_boundary_guard"] and MXG == 0:
raise ValueError("No x guards found in dataset, cannot select radial guards!")

return slices[name]