Merge pull request #64 from eX-Mech/enh/skip-readnek-vars

Skip variable initialization in readnek

docs/asv_bench/benchmarks/bench_readnek.py

@@ -5,3 +5,9 @@
 class ReadnekSuite(Channel3D):
     def time_read(self):
         readnek(self.file)
+
+    def time_read_skip_geom(self):
+        readnek(self.file, skip_vars=("x", "y", "z"))
+
+    def time_read_skip_ux_uy(self):
+        readnek(self.file, skip_vars=("ux", "uy"))

docs/asv_bench/results/benchmarks.json

@@ -31,5 +31,37 @@
         "version": "4247c82d15af70d066d71b05656858fa8f9088d8df24e2be77eb30b987a05a4b",
         "warmup_time": -1
     },
+    "bench_readnek.ReadnekSuite.time_read_skip_geom": {
+        "code": "class ReadnekSuite:\n    def time_read_skip_geom(self):\n        readnek(self.file, skip_vars=(\"x\", \"y\", \"z\"))\n\nclass Channel3D:\n    def setup(self):\n        self.file = test_data / \"nek\" / \"channel3D_0.f00001\"",
+        "min_run_count": 2,
+        "name": "bench_readnek.ReadnekSuite.time_read_skip_geom",
+        "number": 0,
+        "param_names": [],
+        "params": [],
+        "repeat": 0,
+        "rounds": 2,
+        "sample_time": 0.01,
+        "timeout": 60.0,
+        "type": "time",
+        "unit": "seconds",
+        "version": "648b33626bd1d331af4b0fee647b9027cc35a0fd5f0c607992ab741bbdd04a2f",
+        "warmup_time": -1
+    },
+    "bench_readnek.ReadnekSuite.time_read_skip_ux_uy": {
+        "code": "class ReadnekSuite:\n    def time_read_skip_ux_uy(self):\n        readnek(self.file, skip_vars=(\"ux\", \"uy\"))\n\nclass Channel3D:\n    def setup(self):\n        self.file = test_data / \"nek\" / \"channel3D_0.f00001\"",
+        "min_run_count": 2,
+        "name": "bench_readnek.ReadnekSuite.time_read_skip_ux_uy",
+        "number": 0,
+        "param_names": [],
+        "params": [],
+        "repeat": 0,
+        "rounds": 2,
+        "sample_time": 0.01,
+        "timeout": 60.0,
+        "type": "time",
+        "unit": "seconds",
+        "version": "0bb71373a40c99ef4da0c0c9a421ad39ec214e15dd62c479e00faefb798856cb",
+        "warmup_time": -1
+    },
     "version": 2
 }

docs/asv_bench/results/tuxpulse/427ce267-virtualenv-py3.9-dask[array]-numpy-pandas.json

@@ -0,0 +1 @@
+{"commit_hash": "427ce2674ef16a2e27c575279fbcfbfe1397d90a", "env_name": "virtualenv-py3.9-dask[array]-numpy-pandas", "date": 1666220678000, "params": {"arch": "x86_64", "cpu": "AMD Ryzen 7 4800H with Radeon Graphics", "machine": "tuxpulse", "num_cpu": "16", "os": "Linux 5.15.0-10033-tuxedo", "ram": "16GB", "python": "3.9", "numpy": "", "dask[array]": "", "pandas": ""}, "python": "3.9", "requirements": {"numpy": "", "dask[array]": "", "pandas": ""}, "env_vars": {}, "result_columns": ["result", "params", "version", "started_at", "duration", "stats_ci_99_a", "stats_ci_99_b", "stats_q_25", "stats_q_75", "stats_number", "stats_repeat", "samples", "profile"], "results": {"bench_open_dataset.OpenDatasetSuite.time_read": [[0.9001364415000239], [], "d8736c9050c8133f745cc20f6d35fcbac47a834c114c581fff6c8b4123f99d6a", 1666549141800, 10.989, [0.89589], [0.92551], [0.89636], [0.907], [1], [10]], "bench_readnek.ReadnekSuite.time_read": [[0.028870083000015256], [], "4247c82d15af70d066d71b05656858fa8f9088d8df24e2be77eb30b987a05a4b", 1666549147321, 0.61026, [0.028801], [0.02897], [0.028834], [0.028903], [1], [10]], "bench_readnek.ReadnekSuite.time_read_skip_geom": [[0.0232559510000101], [], "648b33626bd1d331af4b0fee647b9027cc35a0fd5f0c607992ab741bbdd04a2f", 1666549147626, 0.50942, [0.023017], [0.02426], [0.023207], [0.023315], [1], [10]], "bench_readnek.ReadnekSuite.time_read_skip_ux_uy": [[0.026239643499991416], [], "0bb71373a40c99ef4da0c0c9a421ad39ec214e15dd62c479e00faefb798856cb", 1666549147882, 0.61017, [0.025937], [0.030505], [0.02605], [0.0264], [1], [10]]}, "durations": {"<build>": 4.768466949462891}, "version": 2}

docs/asv_bench/results/tuxpulse/427ce267-virtualenv-py3.9-numpy-pandas.json

@@ -0,0 +1 @@
+{"commit_hash": "427ce2674ef16a2e27c575279fbcfbfe1397d90a", "env_name": "virtualenv-py3.9-numpy-pandas", "date": 1666220678000, "params": {"arch": "x86_64", "cpu": "AMD Ryzen 7 4800H with Radeon Graphics", "machine": "tuxpulse", "num_cpu": "16", "os": "Linux 5.15.0-10033-tuxedo", "ram": "16GB", "python": "3.9", "numpy": "", "pandas": ""}, "python": "3.9", "requirements": {"numpy": "", "pandas": ""}, "env_vars": {}, "result_columns": ["result", "params", "version", "started_at", "duration", "stats_ci_99_a", "stats_ci_99_b", "stats_q_25", "stats_q_75", "stats_number", "stats_repeat", "samples", "profile"], "results": {"bench_open_dataset.OpenDatasetSuite.time_read": [[0.9033288629999561], [], "d8736c9050c8133f745cc20f6d35fcbac47a834c114c581fff6c8b4123f99d6a", 1666549159773, 10.989, [0.90068], [0.90947], [0.90168], [0.90602], [1], [10]], "bench_readnek.ReadnekSuite.time_read": [[0.028691123999976753], [], "4247c82d15af70d066d71b05656858fa8f9088d8df24e2be77eb30b987a05a4b", 1666549165244, 0.60858, [0.028513], [0.028834], [0.028584], [0.028764], [1], [10]], "bench_readnek.ReadnekSuite.time_read_skip_geom": [[0.0230485534999616], [], "648b33626bd1d331af4b0fee647b9027cc35a0fd5f0c607992ab741bbdd04a2f", 1666549165549, 0.50986, [0.022923], [0.023391], [0.022989], [0.023155], [1], [10]], "bench_readnek.ReadnekSuite.time_read_skip_ux_uy": [[0.026296040499971696], [], "0bb71373a40c99ef4da0c0c9a421ad39ec214e15dd62c479e00faefb798856cb", 1666549165804, 0.60926, [0.026243], [0.026416], [0.026271], [0.026343], [1], [10]]}, "durations": {"<build>": 4.7109808921813965}, "version": 2}

docs/asv_bench/results/tuxpulse/b11e076b-virtualenv-py3.9-dask[array]-numpy-pandas.json

@@ -0,0 +1 @@
+{"commit_hash": "b11e076b96c0fb7cbf73981bff0beaea93f22c3a", "env_name": "virtualenv-py3.9-dask[array]-numpy-pandas", "date": 1665785574000, "params": {"arch": "x86_64", "cpu": "AMD Ryzen 7 4800H with Radeon Graphics", "machine": "tuxpulse", "num_cpu": "16", "os": "Linux 5.15.0-10033-tuxedo", "ram": "16GB", "python": "3.9", "numpy": "", "dask[array]": "", "pandas": ""}, "python": "3.9", "requirements": {"numpy": "", "dask[array]": "", "pandas": ""}, "env_vars": {}, "result_columns": ["result", "params", "version", "started_at", "duration", "stats_ci_99_a", "stats_ci_99_b", "stats_q_25", "stats_q_75", "stats_number", "stats_repeat", "samples", "profile"], "results": {"bench_open_dataset.OpenDatasetSuite.time_read": [[0.9131130614999847], [], "d8736c9050c8133f745cc20f6d35fcbac47a834c114c581fff6c8b4123f99d6a", 1666549060971, 11.089, [0.89895], [0.92779], [0.90003], [0.92572], [1], [10]], "bench_readnek.ReadnekSuite.time_read": [[0.029030419500031712], [], "4247c82d15af70d066d71b05656858fa8f9088d8df24e2be77eb30b987a05a4b", 1666549066592, 0.61056, [0.028948], [0.029107], [0.029005], [0.029042], [1], [10]], "bench_readnek.ReadnekSuite.time_read_skip_geom": [null, [], "648b33626bd1d331af4b0fee647b9027cc35a0fd5f0c607992ab741bbdd04a2f", 1666549060862, 0.054105], "bench_readnek.ReadnekSuite.time_read_skip_ux_uy": [null, [], "0bb71373a40c99ef4da0c0c9a421ad39ec214e15dd62c479e00faefb798856cb", 1666549060916, 0.054584]}, "durations": {}, "version": 2}

docs/asv_bench/results/tuxpulse/b11e076b-virtualenv-py3.9-numpy-pandas.json

@@ -0,0 +1 @@
+{"commit_hash": "b11e076b96c0fb7cbf73981bff0beaea93f22c3a", "env_name": "virtualenv-py3.9-numpy-pandas", "date": 1665785574000, "params": {"arch": "x86_64", "cpu": "AMD Ryzen 7 4800H with Radeon Graphics", "machine": "tuxpulse", "num_cpu": "16", "os": "Linux 5.15.0-10033-tuxedo", "ram": "16GB", "python": "3.9", "numpy": "", "pandas": ""}, "python": "3.9", "requirements": {"numpy": "", "pandas": ""}, "env_vars": {}, "result_columns": ["result", "params", "version", "started_at", "duration", "stats_ci_99_a", "stats_ci_99_b", "stats_q_25", "stats_q_75", "stats_number", "stats_repeat", "samples", "profile"], "results": {"bench_open_dataset.OpenDatasetSuite.time_read": [[0.8863102875000095], [], "d8736c9050c8133f745cc20f6d35fcbac47a834c114c581fff6c8b4123f99d6a", 1666549073462, 10.788, [0.88111], [0.91047], [0.88402], [0.89598], [1], [10]], "bench_readnek.ReadnekSuite.time_read": [[0.0284666934999791], [], "4247c82d15af70d066d71b05656858fa8f9088d8df24e2be77eb30b987a05a4b", 1666549078831, 0.60931, [0.028417], [0.028558], [0.028445], [0.028507], [1], [10]], "bench_readnek.ReadnekSuite.time_read_skip_geom": [null, [], "648b33626bd1d331af4b0fee647b9027cc35a0fd5f0c607992ab741bbdd04a2f", 1666549073355, 0.05351], "bench_readnek.ReadnekSuite.time_read_skip_ux_uy": [null, [], "0bb71373a40c99ef4da0c0c9a421ad39ec214e15dd62c479e00faefb798856cb", 1666549073408, 0.053331]}, "durations": {}, "version": 2}

pymech/neksuite/field.py

@@ -175,7 +175,7 @@ def read_header(path_or_file_obj: Union[PathLike, BinaryIO]) -> Header:
 
 
 # ==============================================================================
-def readnek(fname, dtype="float64"):
+def readnek(fname, dtype="float64", skip_vars=()):
     """A function for reading binary data from the nek5000 binary format
 
     Parameters
@@ -184,6 +184,12 @@ def readnek(fname, dtype="float64"):
         File name
     dtype : str or type
         Floating point data type. See also :class:`pymech.core.Elem`.
+    skip_vars: tuple[str]
+        Variables to skip. Valid values to skip are ``("x", "y", "z", "ux",
+        "uy", "uz", "pressure", "temperature", "s01", "s02", ...)``.  It also
+        accept some extra values ``("vx", "vy", "vz", "p", "t")``.  If empty
+        (default), it reads all variables available in the file.
+
     """
     #
     try:
@@ -234,51 +240,105 @@ def readnek(fname, dtype="float64"):
     elif emode == ">":
         data.endian = "big"
 
+    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(h.nb_pts_elem * h.wdsz)
+        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_elements(nb_elements=1):
+        infile.seek(bytes_elem * nb_elements, os.SEEK_CUR)
+
     # read geometry
-    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)
-    #
+    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))
+    if nb_vars:
+        if all(skip_condition):
+            skip_elements(h.nb_elems * nb_vars)
+        else:
+            for iel in elmap:
+                el = data.elem[iel - 1]
+                for idim in range(nb_vars):
+                    if skip_condition[idim]:
+                        skip_elements()
+                    else:
+                        read_file_into_data(el.pos, idim)
+
     # read velocity
-    for iel in elmap:
-        el = data.elem[iel - 1]
-        for idim in range(h.nb_vars[1]):  # if 0, velocity is not read
-            read_file_into_data(el.vel, idim)
+    velocity_vars1 = "ux", "uy", "uz"
+    velocity_vars2 = "vx", "vy", "vz"
+    nb_vars = h.nb_vars[1]
+    skip_condition1 = tuple(
+        velocity_vars1[idim] in skip_vars for idim in range(nb_vars)
+    )
+    skip_condition2 = tuple(
+        velocity_vars2[idim] in skip_vars for idim in range(nb_vars)
+    )
+
+    if nb_vars:
+        if all(skip_condition1) or all(skip_condition2):
+            skip_elements(h.nb_elems * nb_vars)
+        else:
+            for iel in elmap:
+                el = data.elem[iel - 1]
+                for idim in range(nb_vars):
+                    if skip_condition1[idim] or skip_condition2[idim]:
+                        skip_elements()
+                    else:
+                        read_file_into_data(el.vel, idim)
+
     #
     # 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

@@ -4,6 +4,7 @@
 
 import numpy as np
 from numpy import testing as npt
+import pytest
 
 from pymech.log import logger
 
@@ -139,6 +140,48 @@ def test_writenek_scalars(test_data_dir, tmpdir):
     npt.assert_array_equal(field.lims.scal, fieldw.lims.scal)
 
 
+def test_readnek_skip_vars(test_data_dir):
+    import pymech.neksuite as ns
+
+    fname = f"{test_data_dir}/nek/channel3D_0.f00001"
+
+    field_all = ns.readnek(fname)
+    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)
+        with pytest.raises(AssertionError):
+            npt.assert_array_equal(elem.pos, elem_skip_geom.pos)
+
+        with pytest.raises(AssertionError):
+            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