Implement skip_vars for pressure, temperature and scalars

pymech/neksuite/field.py

@@ -259,7 +259,6 @@ def skip_elements(nb_elements=1):
     geometry_vars = "x", "y", "z"
     nb_vars = h.nb_vars[0]
     skip_condition = tuple(geometry_vars[idim] in skip_vars for idim in range(nb_vars))
-    # breakpoint()
     if nb_vars:
         if all(skip_condition):  # :nb_vars
             skip_elements(h.nb_elems * nb_vars)
@@ -297,27 +296,49 @@ def skip_elements(nb_elements=1):
 
     #
     # read pressure
-    for iel in elmap:
-        el = data.elem[iel - 1]
-        for ivar in range(h.nb_vars[2]):  # if 0, pressure is not read
-            read_file_into_data(el.pres, ivar)
+    nb_vars = h.nb_vars[2]
+    skip_condition = any({"p", "pressure"}.intersection(skip_vars))
+    if nb_vars:
+        if skip_condition:
+            skip_elements(h.nb_elems * nb_vars)
+        else:
+            for iel in elmap:
+                el = data.elem[iel - 1]
+                for ivar in range(nb_vars):
+                    read_file_into_data(el.pres, ivar)
+
     #
     # read temperature
-    for iel in elmap:
-        el = data.elem[iel - 1]
-        for ivar in range(h.nb_vars[3]):  # if 0, temperature is not read
-            read_file_into_data(el.temp, ivar)
+    nb_vars = h.nb_vars[3]
+    skip_condition = any({"t", "temperature"}.intersection(skip_vars))
+    if nb_vars:
+        if skip_condition:
+            skip_elements(h.nb_elems * nb_vars)
+        else:
+            for iel in elmap:
+                el = data.elem[iel - 1]
+                for ivar in range(nb_vars):
+                    read_file_into_data(el.temp, ivar)
     #
     # read scalar fields
     #
-    # NOTE: This is not a bug!
-    # Unlike other variables, scalars are in the outer loop and elements
-    # are in the inner loop
-    #
-    for ivar in range(h.nb_vars[4]):  # if 0, scalars are not read
-        for iel in elmap:
-            el = data.elem[iel - 1]
-            read_file_into_data(el.scal, ivar)
+    nb_vars = h.nb_vars[4]
+    scalar_vars = tuple(f"s{i:02d}" for i in range(1, nb_vars + 1))
+    skip_condition = tuple(scalar_vars[ivar] in skip_vars for ivar in range(nb_vars))
+    if nb_vars:
+        if all(skip_condition):
+            skip_elements(h.nb_elems * nb_vars)
+        else:
+            # NOTE: This is not a bug!
+            # Unlike other variables, scalars are in the outer loop and elements
+            # are in the inner loop
+            for ivar in range(nb_vars):
+                if skip_condition[ivar]:
+                    skip_elements(h.nb_elems)
+                else:
+                    for iel in elmap:
+                        el = data.elem[iel - 1]
+                        read_file_into_data(el.scal, ivar)
     #
     #
     # close file

tests/test_all.py

@@ -165,6 +165,23 @@ def test_readnek_skip_vars(test_data_dir):
             npt.assert_array_equal(elem.vel, elem_skip_ux_uy.vel)
 
 
+def test_readnek_skip_scalars(test_data_dir):
+    import pymech.neksuite as ns
+
+    fname = f"{test_data_dir}/nek/stsabl0.f00001"
+    field_all = ns.readnek(fname)
+    field_skip = ns.readnek(fname, skip_vars=("s02, s04"))
+    for elem, elem_skip in zip(field_all.elem, field_skip.elem):
+        npt.assert_array_equal(elem.scal[0], elem_skip.scal[0])
+        with pytest.raises(AssertionError):
+            npt.assert_array_equal(elem.scal[1], elem_skip.scal[1])  # s02
+
+        npt.assert_array_equal(elem.scal[2], elem_skip.scal[2])
+
+        with pytest.raises(AssertionError):
+            npt.assert_array_equal(elem.scal[3], elem_skip.scal[3])  # s04
+
+
 def test_readrea(test_data_dir):
     import pymech.neksuite as ns