Skip to content
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.

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

feat: add ikle2/ikle3 as variable #40

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
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
12 changes: 6 additions & 6 deletions tests/io_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,9 +25,9 @@
def write_netcdf(ds, nc_out):
# Remove dict and multi-dimensional arrays not supported in netCDF
del ds.attrs["variables"]
del ds.attrs["ikle2"]
del ds["face"]
try:
del ds.attrs["ikle3"]
del ds.attrs["volume"]
except KeyError:
pass
# Write netCDF file
Expand Down Expand Up @@ -55,11 +55,11 @@ def test_open_dataset(slf_in):
assert ds.attrs["endian"] == ">"
assert ds.attrs["date_start"] == (1900, 1, 1, 0, 0, 0)
assert "ipobo" in ds.attrs
assert "ikle2" in ds.attrs
assert "face" in ds.variables
if "r3d_" in slf_in:
assert "ikle3" in ds.attrs
assert "volume" in ds.variables
else:
assert "ikle3" not in ds.attrs
assert "volume" not in ds.variables


@TIDAL_FLATS
Expand Down Expand Up @@ -155,9 +155,9 @@ def test_from_scratch(tmp_path):
"x": ("node", x),
"y": ("node", y),
"time": pd.date_range("2020-01-01", periods=10),
"face": (("face_node_connectivity", "face_dimension"), ikle),
},
)
ds.attrs["ikle2"] = ikle

# Writing to a SELAFIN file
ds.selafin.write(slf_out)
Expand Down
52 changes: 39 additions & 13 deletions xarray_selafin/xarray_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,11 +66,15 @@ def write_serafin(fout, ds):
first_time = ds.time
else:
first_time = ds.time[0]
first_date_str = first_time.values.astype(str) # "1900-01-01T00:00:00.000000000"
first_date_str = first_time.values.astype(
str
) # "1900-01-01T00:00:00.000000000"
first_date_str = first_date_str.rstrip("0") + "0" # "1900-01-01T00:00:00.0"
try:
date = datetime.strptime(first_date_str, '%Y-%m-%dT%H:%M:%S.%f')
slf_header.date = attrgetter("year", "month", "day", "hour", "minute", "second")(date)
date = datetime.strptime(first_date_str, "%Y-%m-%dT%H:%M:%S.%f")
slf_header.date = attrgetter(
"year", "month", "day", "hour", "minute", "second"
)(date)
except ValueError:
slf_header.date = DEFAULT_DATE_START

Expand All @@ -94,12 +98,12 @@ def write_serafin(fout, ds):
is_2d = False
nplan = len(ds.plan)
slf_header.nb_nodes_per_elem = 6
slf_header.nb_elements = len(ds.attrs["ikle2"]) * (nplan - 1)
slf_header.nb_elements = len(ds["face"]) * (nplan - 1)
else: # 2D
is_2d = True
nplan = 1 # just to do a multiplication
slf_header.nb_nodes_per_elem = ds.attrs["ikle2"].shape[1]
slf_header.nb_elements = len(ds.attrs["ikle2"])
slf_header.nb_nodes_per_elem = ds["face"].shape[1]
slf_header.nb_elements = len(ds["face"])

slf_header.nb_nodes = ds.sizes["node"] * nplan
slf_header.nb_nodes_2d = ds.sizes["node"]
Expand All @@ -114,15 +118,17 @@ def write_serafin(fout, ds):
slf_header.mesh_origin = (0, 0) # Should be integers
slf_header.x_stored = x - slf_header.mesh_origin[0]
slf_header.y_stored = y - slf_header.mesh_origin[1]
slf_header.ikle_2d = ds.attrs["ikle2"]
slf_header.ikle_2d = np.array(ds["face"])
if is_2d:
slf_header.ikle = slf_header.ikle_2d.flatten()
else:
try:
slf_header.ikle = ds.attrs["ikle3"]
slf_header.ikle = ds["volume"]
except KeyError:
# Rebuild IKLE from 2D
slf_header.ikle = slf_header.compute_ikle(len(ds.plan), slf_header.nb_nodes_per_elem)
slf_header.ikle = slf_header.compute_ikle(
len(ds.plan), slf_header.nb_nodes_per_elem
)

try:
slf_header.ipobo = ds.attrs["ipobo"]
Expand Down Expand Up @@ -308,28 +314,48 @@ def open_dataset(
data[time_index, :, :] = values.reshape(npoin2, nplan)
data_vars[var] = xr.Variable(dims=dims, data=data)

# include UGRID conventions: https://ugrid-conventions.github.io/ugrid-conventions/conformance/#R102
# topology_dimension = 0 (points)
# topology_dimension = 1 (edge) > does not exist for Selafin files
# topology_dimension = 2 (face)
# topology_dimension = 3 (volume)

coords = {
"x": ("node", x[:npoin2]),
"y": ("node", y[:npoin2]),
"time": times,
"face": (("face_node_connectivity", "face_dimension"), slf.header.ikle_2d)
# Consider how to include IPOBO (with node and plan dimensions?) if it's essential for your analysis
}
if not is_2d: # if 3D
ikle3 = np.reshape(slf.header.ikle, (slf.header.nb_elements, ndp3))
coords["volume"] = (("volume_node_connectivity", "volume_dimension"), ikle3)

# build node connectivity:
coords["boundary"] = (
"boundary_node_connectivity",
list(filter(lambda num: num != 0, slf.header.ipobo)),
)

ds = xr.Dataset(data_vars=data_vars, coords=coords)

if len(ds.face_dimension) == 1:
ds.attrs["topology_dimension"] = 0
elif len(ds.face_dimension) == 3:
ds.attrs["topology_dimension"] = 2
else: # len(face_dimension) > 3
ds.attrs["topology_dimension"] = 3

ds.attrs["title"] = slf.header.title.decode(Serafin.SLF_EIT).strip()
ds.attrs["language"] = slf.header.language
ds.attrs["float_size"] = slf.header.float_size
ds.attrs["endian"] = slf.header.endian
ds.attrs["params"] = slf.header.params
ds.attrs["ipobo"] = slf.header.ipobo
ds.attrs["ikle2"] = slf.header.ikle_2d
if not is_2d:
ds.attrs["ikle3"] = np.reshape(slf.header.ikle, (slf.header.nb_elements, ndp3))
ds.attrs["variables"] = {
var_ID: (
name.decode(Serafin.SLF_EIT).rstrip(),
unit.decode(Serafin.SLF_EIT).rstrip()
unit.decode(Serafin.SLF_EIT).rstrip(),
)
for var_ID, name, unit in slf.header.iter_on_all_variables()
}
Expand Down
Loading