Fix skip_data, renamed as skip_elements

pymech/neksuite/field.py

@@ -238,30 +238,35 @@ def readnek(fname, dtype="float64", skip_vars=()):
     elif emode == ">":
         data.endian = "big"
 
-    data_chunk_size = h.nb_pts_elem * h.wdsz
+    bytes_elem = h.nb_pts_elem * h.wdsz
 
     def read_file_into_data(data_var, index_var):
         """Read binary file into an array attribute of ``data.elem``"""
-        fi = infile.read(data_chunk_size)
+        fi = infile.read(bytes_elem)
         fi = np.frombuffer(fi, dtype=emode + h.realtype, count=h.nb_pts_elem)
 
         # Replace elem array in-place with
         # array read from file after reshaping as
         elem_shape = h.orders[::-1]  # lz, ly, lx
         data_var[index_var, ...] = fi.reshape(elem_shape)
 
-    def skip_data(multiplier=1):
-        infile.seek(data_chunk_size * multiplier)
+    def skip_elements(nb_elements=1):
+        infile.seek(bytes_elem * nb_elements, os.SEEK_CUR)
 
     #
     # read geometry
-    if {"x", "y", "z"}.issubset(skip_vars):
-        skip_data(h.nb_elems * h.nb_vars[0])
+    geometry_vars = "x", "y", "z"
+    if set(geometry_vars).issubset(skip_vars):
+        skip_elements(h.nb_elems * h.nb_vars[0])
     else:
         for iel in elmap:
             el = data.elem[iel - 1]
             for idim in range(h.nb_vars[0]):  # if 0, geometry is not read
-                read_file_into_data(el.pos, idim)
+                if geometry_vars[idim] in skip_vars:
+                    skip_elements()
+                else:
+                    read_file_into_data(el.pos, idim)
+    # TODO: skip velocity, pressure and scalars
     #
     # read velocity
     for iel in elmap:

tests/test_all.py

@@ -144,11 +144,18 @@ def test_readnek_skip_vars(test_data_dir):
 
     fname = f"{test_data_dir}/nek/channel3D_0.f00001"
     field_all = ns.readnek(fname)
-    field_skip = ns.readnek(fname, skip_vars=("x", "y", "z"))
-
-    for elem, elemw in zip(field_all.elem, field_skip.elem):
-        npt.assert_array_equal(elem.vel, elemw.vel)
-        npt.assert_array_equal(elem.pres, elemw.pres)
+    field_skip_geom = ns.readnek(fname, skip_vars=("x", "y", "z"))
+    field_skip_ux_uy = ns.readnek(fname, skip_vars=("ux", "uy"))
+
+    for elem, elem_skip_geom, elem_skip_ux_uy in zip(
+        field_all.elem, field_skip_geom.elem, field_skip_ux_uy.elem
+    ):
+        npt.assert_array_equal(elem.vel, elem_skip_geom.vel)
+        npt.assert_array_equal(elem.pres, elem_skip_geom.pres)
+        npt.assert_array_equal(elem.scal, elem_skip_geom.scal)
+        npt.assert_array_equal(elem.pos, elem_skip_ux_uy.pos)
+        npt.assert_array_equal(elem.pres, elem_skip_ux_uy.pres)
+        npt.assert_array_equal(elem.scal, elem_skip_ux_uy.scal)
 
 
 def test_readrea(test_data_dir):