-
Notifications
You must be signed in to change notification settings - Fork 10
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge branch 'master' into fix-gridfile-loading
Conflicts: xbout/load.py
- Loading branch information
Showing
16 changed files
with
682 additions
and
305 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Empty file.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,93 @@ | ||
import numpy as np | ||
import xarray as xr | ||
|
||
from .triangulate import Triangulate | ||
|
||
|
||
def da_with_cherab_grid(da): | ||
""" | ||
Convert an BOUT++ DataArray to a format that Cherab can use: | ||
- A 'cell_number' coordinate | ||
- A 'cherab_grid` attribute | ||
The cell_number coordinate enables the DataArray to be sliced | ||
before input to Cherab. | ||
Parameters | ||
---------- | ||
ds : xarray.DataArray | ||
Returns | ||
------- | ||
updated_da | ||
""" | ||
if "cherab_grid" in da.attrs: | ||
# Already has required attribute | ||
return da | ||
|
||
if da.attrs["geometry"] != "toroidal": | ||
raise ValueError("xhermes.plotting.cherab: Geometry must be toroidal") | ||
|
||
if da.sizes["zeta"] != 1: | ||
raise ValueError("xhermes.plotting.cherab: Zeta index must have size 1") | ||
|
||
nx = da.sizes["x"] | ||
ny = da.sizes["theta"] | ||
|
||
# Cell corners | ||
rm = np.stack( | ||
( | ||
da.coords["Rxy_upper_right_corners"], | ||
da.coords["Rxy_upper_left_corners"], | ||
da.coords["Rxy_lower_right_corners"], | ||
da.coords["Rxy_lower_left_corners"], | ||
), | ||
axis=-1, | ||
) | ||
zm = np.stack( | ||
( | ||
da.coords["Zxy_upper_right_corners"], | ||
da.coords["Zxy_upper_left_corners"], | ||
da.coords["Zxy_lower_right_corners"], | ||
da.coords["Zxy_lower_left_corners"], | ||
), | ||
axis=-1, | ||
) | ||
|
||
grid = Triangulate(rm, zm) | ||
|
||
# Store the cell number as a coordinate. | ||
# This allows slicing of arrays before passing to Cherab | ||
|
||
# Create a DataArray for the vertices and triangles | ||
cell_number = xr.DataArray( | ||
grid.cell_number, dims=["x", "theta"], name="cell_number" | ||
) | ||
|
||
result = da.assign_coords(cell_number=cell_number) | ||
result.attrs.update(cherab_grid=grid) | ||
return result | ||
|
||
|
||
def ds_with_cherab_grid(ds): | ||
""" | ||
Create an xarray DataSet with a Cherab grid | ||
Parameters | ||
---------- | ||
ds : xarray.Dataset | ||
Returns | ||
------- | ||
updated_ds | ||
""" | ||
|
||
# The same operation works on a DataSet | ||
ds = da_with_cherab_grid(ds) | ||
|
||
# Add the Cherab grid as an attribute to all variables | ||
grid = ds.attrs["cherab_grid"] | ||
for var in ds.data_vars: | ||
ds[var].attrs.update(cherab_grid=grid) | ||
|
||
return ds |
Oops, something went wrong.