Merge pull request #46 from eX-Mech/enh/single-prec

Single precision dtype option

CHANGELOG.md

@@ -39,6 +39,7 @@ Security in case of vulnerabilities.
 - Attributes `lr1` and `var` {class}`pymech.core.Elem` are tuples instead of lists
 - Attributes of {class}`pymech.core.DataLims` are now immutable tuples
 - {func}`pymech.dataset.open_mfdataset` is now a partial function. This change should be fully backwards compatible
+- {func}`pymech.neksuite.readnek` has a `dtype` option to set the floating point data type.
 
 ### Deprecated
 

pymech/core.py

@@ -1,13 +1,30 @@
 """Core data structures for pymech"""
 import copy
 from textwrap import dedent, indent
+import itertools
 from itertools import product
-from functools import reduce
+from functools import reduce, partial
 
 import numpy as np
 from pymech.log import logger
 
 
+"""Repeat N times. Pythonic idiom to use when the iterated value is discarded.
+
+Example
+-------
+Instead of:
+
+>>> [0 for _ in range(10)]
+
+You could use:
+
+>>> [0 for _ in repeat(10)]
+
+"""
+repeat = partial(itertools.repeat, None)
+
+
 # ==============================================================================
 class DataLims:
     """A class containing the extrema of all quantities stored in the mesh
@@ -78,28 +95,40 @@ def _lims_aggregator(self, elem1, elem2):
 # ==============================================================================
 class Elem:
     """A class containing one hexahedral element of Nek5000/SIMSON flow
-    field
+    field.
+
+    Parameters
+    ----------
+    var : iterable
+        Iterable of integers of size 5, indicating how many variables are to be initialized
+    lr1 : iterable
+        Iterable of integers of size 3, defining the shape of an element as ``(lx, ly, lz)``
+    nbc : int
+        Number of boundary conditions
+    dtype : str
+        Floating point data type. Typical values are 'f4' or 'float32' for
+        single precision, 'f8' or 'float64' for double precision
 
     """
 
-    def __init__(self, var, lr1, nbc):
+    def __init__(self, var, lr1, nbc, dtype="float64"):
         #                    x,y,z   lz      ly      lx
-        self.pos = np.zeros((3, lr1[2], lr1[1], lr1[0]))
+        self.pos = np.zeros((3, lr1[2], lr1[1], lr1[0]), dtype=dtype)
         #                    one per edge
-        self.curv = np.zeros((12, 5))
+        self.curv = np.zeros((12, 5), dtype=dtype)
         #             curvature type
-        self.ccurv = ["" for i in range(12)]
+        self.ccurv = ["" for _ in repeat(12)]
         #                    u,v,w   lz      ly      lx
-        self.vel = np.zeros((3, lr1[2], lr1[1], lr1[0]))
+        self.vel = np.zeros((3, lr1[2], lr1[1], lr1[0]), dtype=dtype)
         #                    p       lz      ly      lx
-        self.pres = np.zeros((var[2], lr1[2], lr1[1], lr1[0]))
+        self.pres = np.zeros((var[2], lr1[2], lr1[1], lr1[0]), dtype=dtype)
         #                    T       lz      ly      lx
-        self.temp = np.zeros((var[3], lr1[2], lr1[1], lr1[0]))
+        self.temp = np.zeros((var[3], lr1[2], lr1[1], lr1[0]), dtype=dtype)
         #                    s_i     lz      ly      lx
-        self.scal = np.zeros((var[4], lr1[2], lr1[1], lr1[0]))
+        self.scal = np.zeros((var[4], lr1[2], lr1[1], lr1[0]), dtype=dtype)
         #                    list of 8 parameters, one per face
         #                    one column for velocity, one for temperature, and one for each scalar
-        self.bcs = np.zeros((nbc, 6), dtype="U3, i4, i4, f8, f8, f8, f8, f8")
+        self.bcs = np.zeros((nbc, 6), dtype="U3, i4, i4" + f", {dtype}" * 5)
 
     def __repr__(self):
         message = f"<elem centered at {self.centroid}>"
@@ -196,7 +225,7 @@ def face_center(self, i):
 class HexaData:
     """A class containing data related to a hexahedral mesh"""
 
-    def __init__(self, ndim, nel, lr1, var, nbc=0):
+    def __init__(self, ndim, nel, lr1, var, nbc=0, dtype="float64"):
         self.ndim = ndim
         self.nel = nel
         self.ncurv = []
@@ -207,7 +236,11 @@ def __init__(self, ndim, nel, lr1, var, nbc=0):
         self.istep = []
         self.wdsz = []
         self.endian = []
-        self.elem = [Elem(var, lr1, nbc) for i in range(nel)]
+        if isinstance(dtype, type):
+            # For example np.float64 -> "float64"
+            dtype = dtype.__name__
+
+        self.elem = [Elem(var, lr1, nbc, dtype) for _ in repeat(nel)]
         self.elmap = np.linspace(1, nel, nel, dtype=np.int32)
 
     def __repr__(self):

pymech/neksuite.py

@@ -150,13 +150,15 @@ def read_header(fp: io.BufferedReader) -> Header:
 
 
 # ==============================================================================
-def readnek(fname):
+def readnek(fname, dtype="float64"):
     """A function for reading binary data from the nek5000 binary format
 
     Parameters
     ----------
     fname : str
-            file name
+        File name
+    dtype : str or type
+        Floating point data type. See also :class:`pymech.core.Elem`.
     """
     #
     try:
@@ -190,14 +192,14 @@ def readnek(fname):
     #
     # read element map for the file
     elmap = infile.read(4 * h.nb_elems_file)
-    elmap = list(struct.unpack(emode + h.nb_elems_file * "i", elmap))
+    elmap = struct.unpack(emode + h.nb_elems_file * "i", elmap)
     #
     # ---------------------------------------------------------------------------
     # READ DATA
     # ---------------------------------------------------------------------------
     #
     # initialize data structure
-    data = HexaData(h.nb_dims, h.nb_elems, h.orders, h.nb_vars, 0)
+    data = HexaData(h.nb_dims, h.nb_elems, h.orders, h.nb_vars, 0, dtype)
     data.time = h.time
     data.istep = h.istep
     data.wdsz = h.wdsz

tests/test_all.py

@@ -2,6 +2,7 @@
 import time
 from textwrap import dedent
 
+import numpy as np
 from numpy import testing as npt
 
 from pymech.log import logger
@@ -114,11 +115,13 @@ def test_readnek_scalars(test_data_dir):
     import pymech.neksuite as ns
 
     # 2D statistics file
+    dtype = np.float32
     fname = f"{test_data_dir}/nek/stsabl0.f00001"
-    field = ns.readnek(fname)
+    field = ns.readnek(fname, dtype)
 
     ux_min, ux_max = field.lims.scal[0]
     assert math.isclose(ux_max, 5.3, abs_tol=0.1)
+    assert field.elem[0].scal.dtype == np.dtype(dtype)
 
 
 def test_writenek_scalars(test_data_dir, tmpdir):
@@ -409,7 +412,6 @@ def test_delete_internal_bcs(test_data_dir):
 def test_extrude(test_data_dir):
     import pymech.neksuite as ns
     import pymech.meshtools as mt
-    import numpy as np
 
     fname = f"{test_data_dir}/nek/2D_section_R360.re2"
     nz = 4
@@ -426,7 +428,6 @@ def test_extrude(test_data_dir):
 
 
 def test_extrude_refine(test_data_dir):
-    import numpy as np
     import pymech.neksuite as ns
     import pymech.meshtools as mt
     from itertools import product