Exodus writer does not write `num_elem` dimension

[philipc2]

uxarray/uxarray/io/_exodus.py

Lines 109 to 293 in ff253ed

	def _encode_exodus(ds, outfile=None):
	"""Encodes an Exodus file.
	Parameters
	----------
	ds : xarray.Dataset, required
	Dataset to be encoded to exodus file.
	outfile : string, required
	Name of output file to be added as metadata into the output
	dataset
	Returns
	-------
	exo_ds : xarray.Dataset
	Dataset encoded as exodus file.
	"""
	# Note this is 1-based unlike native Mesh2 construct
	exo_ds = xr.Dataset()
	now = datetime.now()
	date = now.strftime("%Y:%m:%d")
	time = now.strftime("%H:%M:%S")
	fp_word = INT_DTYPE(8)
	exo_version = np.float32(5.0)
	api_version = np.float32(5.0)
	exo_ds.attrs = {
	"api_version": api_version,
	"version": exo_version,
	"floating_point_word_size": fp_word,
	"file_size": 0,
	}
	if outfile:
	path = PurePath(outfile)
	out_filename = path.name
	title = "uxarray(" + str(out_filename) + ")" + date + ": " + time
	exo_ds.attrs = {
	"api_version": api_version,
	"version": exo_version,
	"floating_point_word_size": fp_word,
	"file_size": 0,
	"title": title,
	}
	exo_ds["time_whole"] = xr.DataArray(data=[], dims=["time_step"])
	# qa_records
	# version identifier of the application code: https://gsjaardema.github.io/seacas-docs/exodusII-new.pdf page: 12
	ux_exodus_version = 1.0
	qa_records = [["uxarray"], [ux_exodus_version], [date], [time]]
	exo_ds["qa_records"] = xr.DataArray(
	data=xr.DataArray(np.array(qa_records, dtype="str")),
	dims=["four", "num_qa_rec"],
	)
	# Set the dim to 3 as we will always have x/y/z for cartesian grid
	# Note: Don't get orig dimension from Mesh2 attribute topology dimension
	if "node_x" not in ds:
	x, y, z = _lonlat_rad_to_xyz(ds["node_lon"].values, ds["node_lat"].values)
	c_data = xr.DataArray([x, y, z])
	else:
	c_data = xr.DataArray(
	[
	ds["node_x"].data.tolist(),
	ds["node_y"].data.tolist(),
	ds["node_z"].data.tolist(),
	]
	)
	exo_ds["coord"] = xr.DataArray(data=c_data, dims=["num_dim", "num_nodes"])
	# process face nodes, this array holds num faces at corresponding location
	# eg num_el_all_blks = [0, 0, 6, 12] signifies 6 TRI and 12 SHELL elements
	num_el_all_blks = np.zeros(ds["n_max_face_nodes"].size, "i8")
	# this list stores connectivity without filling
	conn_nofill = []
	# store the number of faces in an array
	for row in ds["face_node_connectivity"].astype(INT_DTYPE).data:
	# find out -1 in each row, this indicates lower than max face nodes
	arr = np.where(row == -1)
	# arr[0].size returns the location of first -1 in the conn list
	# if > 0, arr[0][0] is the num_nodes forming the face
	if arr[0].size > 0:
	# increment the number of faces at the corresponding location
	num_el_all_blks[arr[0][0] - 1] += 1
	# append without -1s eg. [1, 2, 3, -1] to [1, 2, 3]
	# convert to list (for sorting later)
	row = row[: (arr[0][0])].tolist()
	list_node = list(map(int, row))
	conn_nofill.append(list_node)
	elif arr[0].size == 0:
	# increment the number of faces for this "n_max_face_nodes" face
	num_el_all_blks[ds["n_max_face_nodes"].size - 1] += 1
	# get integer list nodes
	list_node = list(map(int, row.tolist()))
	conn_nofill.append(list_node)
	else:
	raise RuntimeError(
	"num nodes in conn array is greater than n_max_face_nodes. Abort!"
	)
	# get number of blks found
	num_blks = np.count_nonzero(num_el_all_blks)
	# sort connectivity by size, lower dim faces first
	conn_nofill.sort(key=len)
	# get index of blocks found
	nonzero_el_index_blks = np.nonzero(num_el_all_blks)
	# break face_node_connectivity into blks
	start = 0
	for blk in range(num_blks):
	blkID = blk + 1
	str_el_in_blk = "num_el_in_blk" + str(blkID)
	str_nod_per_el = "num_nod_per_el" + str(blkID)
	str_att_in_blk = "num_att_in_blk" + str(blkID)
	str_global_id = "global_id" + str(blkID)
	str_edge_type = "edge_type" + str(blkID)
	str_attrib = "attrib" + str(blkID)
	str_connect = "connect" + str(blkID)
	# get element type
	num_nodes = len(conn_nofill[start])
	element_type = _get_element_type(num_nodes)
	# get number of faces for this block
	num_faces = num_el_all_blks[nonzero_el_index_blks[0][blk]]
	# assign Data variables
	# convert list to np.array, sorted list guarantees we have the correct info
	conn_blk = conn_nofill[start : start + num_faces]
	conn_np = np.array([np.array(xi, dtype=INT_DTYPE) for xi in conn_blk])
	exo_ds[str_connect] = xr.DataArray(
	data=xr.DataArray((conn_np[:] + 1)),
	dims=[str_el_in_blk, str_nod_per_el],
	attrs={"elem_type": element_type},
	)
	# edge type
	exo_ds[str_edge_type] = xr.DataArray(
	data=xr.DataArray(np.zeros((num_faces, num_nodes), dtype=INT_DTYPE)),
	dims=[str_el_in_blk, str_nod_per_el],
	)
	# global id
	gid = np.arange(start + 1, start + num_faces + 1, 1)
	exo_ds[str_global_id] = xr.DataArray(data=(gid), dims=[str_el_in_blk])
	# attrib
	# TODO: fix num attr
	num_attr = 1
	exo_ds[str_attrib] = xr.DataArray(
	data=xr.DataArray(np.zeros((num_faces, num_attr), float)),
	dims=[str_el_in_blk, str_att_in_blk],
	)
	start = num_faces
	# blk for loop ends
	# eb_prop1
	prop1_vals = np.arange(1, num_blks + 1, 1)
	exo_ds["eb_prop1"] = xr.DataArray(
	data=prop1_vals, dims=["num_el_blk"], attrs={"name": "ID"}
	)
	# eb_status
	exo_ds["eb_status"] = xr.DataArray(
	data=xr.DataArray(np.ones([num_blks], dtype=INT_DTYPE)), dims=["num_el_blk"]
	)
	# eb_names
	eb_names = np.empty(num_blks, dtype="str")
	exo_ds["eb_names"] = xr.DataArray(data=xr.DataArray(eb_names), dims=["num_el_blk"])
	cnames = ["x", "y", "z"]
	exo_ds["coor_names"] = xr.DataArray(
	data=xr.DataArray(np.array(cnames, dtype="str")), dims=["num_dim"]
	)
	return exo_ds