ENH: Use surfio for irap ascii

Author: eivindjahren

pyproject.toml

@@ -49,6 +49,7 @@ dependencies = [
     "scipy>=1.5",
     "segyio>1.8.0",
     "shapely>=1.6.2",
+    "surfio",
     "tables",
     "typing_extensions",
     "xtgeoviz",

src/xtgeo/surface/_regsurf_export.py

@@ -2,14 +2,14 @@
 
 from __future__ import annotations
 
-import io
 import json
 import struct
 from typing import TYPE_CHECKING
 
 import h5py
 import hdf5plugin
 import numpy as np
+import surfio
 
 from xtgeo import _cxtgeo
 from xtgeo.common.constants import UNDEF_MAP_IRAPA, UNDEF_MAP_IRAPB
@@ -47,113 +47,64 @@
 def export_irap_ascii(self: RegularSurface, mfile: FileWrapper) -> None:
     """Export to Irap RMS ascii format."""
 
-    vals = self.get_values1d(fill_value=UNDEF_MAP_IRAPA, order="F")
+    vals = np.ma.filled(self.values, fill_value=UNDEF_MAP_IRAPA)
 
     yinc = self.yinc * self.yflip
 
     xmax = self.xori + (self.ncol - 1) * self.xinc
     ymax = self.yori + (self.nrow - 1) * yinc
 
-    header = (
-        f"-996 {self.nrow} {self.xinc} {yinc}\n"
-        f"{self.xori} {xmax} {self.yori} {ymax}\n"
-        f"{self.ncol} {self.rotation} {self.xori} {self.yori}\n"
-        "0  0  0  0  0  0  0\n"
+    surf = surfio.IrapSurface(
+        surfio.IrapHeader(
+            ncol=self.ncol,
+            nrow=self.nrow,
+            xori=self.xori,
+            yori=self.yori,
+            xinc=self.xinc,
+            yinc=yinc,
+            xmax=xmax,
+            ymax=ymax,
+            rot=self.rotation,
+            xrot=self.xori,
+            yrot=self.yori,
+        ),
+        values=vals,
     )
 
-    def _optimal_shape(vals, start=9):
-        """Optimal shape for the data.
-
-        It seems by reverse engineering that RMS accepts only 9 or less items per line
-        """
-        # Check if divisible by a number
-        size = len(vals)
-        # Find the nearest factorization
-        for i in range(start, 1, -1):
-            if size % i == 0:
-                return (int(size // i), int(i))  # Ensure integers are returned
-
-        # If we can't find a perfect divisor, return a valid shape
-        return (int(size), 1)
-
-    buffer = io.StringIO()
-
-    np.savetxt(buffer, vals.reshape(_optimal_shape(vals)), fmt="%f", delimiter=" ")
-    data = buffer.getvalue()
-
-    # Combine header and data
-    buf = (header + data).encode("latin1")
-
     if mfile.memstream:
-        mfile.file.write(buf)
+        mfile.file.write(surf.to_ascii_string().encode("latin1"))
     else:
-        with open(mfile.name, "wb") as fout:
-            fout.write(buf)
-
-    del vals
+        surf.to_ascii_file(mfile.name)
 
 
 def export_irap_binary(self: RegularSurface, mfile: FileWrapper) -> None:
     """Export to Irap RMS binary format."""
 
-    vals = self.get_values1d(fill_value=UNDEF_MAP_IRAPB, order="F")
-
     yinc = self.yinc * self.yflip
-
-    header = struct.pack(
-        ">3i6f3i3f10i",  # > means big endian storage
-        32,
-        -996,
-        self.nrow,
-        self.xori,
-        self.xori + self.xinc * (self.ncol - 1),
-        self.yori,
-        self.yori + yinc * (self.nrow - 1),
-        self.xinc,
-        yinc,
-        32,
-        16,
-        self.ncol,
-        self.rotation,
-        self.xori,
-        self.yori,
-        16,
-        28,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        28,
+    xmax = self.xori + (self.ncol - 1) * self.xinc
+    ymax = self.yori + (self.nrow - 1) * yinc
+    values = np.ma.filled(self.values, fill_value=UNDEF_MAP_IRAPB)
+
+    surf = surfio.IrapSurface(
+        surfio.IrapHeader(
+            ncol=self.ncol,
+            nrow=self.nrow,
+            xori=self.xori,
+            yori=self.yori,
+            xinc=self.xinc,
+            yinc=yinc,
+            xmax=xmax,
+            ymax=ymax,
+            rot=self.rotation,
+            xrot=self.xori,
+            yrot=self.yori,
+        ),
+        values=values,
     )
-    inum = self.nrow * self.ncol
-
-    # export to Irap binary in ncol chunks (the only chunk size accepted by RMS)
-    nchunk = self.ncol
-    chunks = [nchunk] * (inum // nchunk)
-    if (inum % nchunk) > 0:
-        chunks.append(inum % nchunk)
-    start = 0
-    data = bytearray(header)
-    chunk_size = nchunk * 4
-
-    # Precompute the struct.pack format for chunk size
-    chunk_size_pack = struct.pack(">i", chunk_size)
-
-    for chunk in chunks:
-        chunk_data = np.array(vals[start : start + chunk], dtype=">f4").tobytes()
-        data.extend(chunk_size_pack)
-        data.extend(chunk_data)
-        data.extend(chunk_size_pack)
-        start += chunk
-
     if mfile.memstream:
-        mfile.file.write(data)
+        mfile.file.write(surf.to_binary_buffer())
     else:
-        with open(mfile.name, "wb") as fout:
-            fout.write(data)
+        surf.to_binary_file(mfile.file)
 
 
 def export_ijxyz_ascii(self: RegularSurface, mfile: FileWrapper) -> None:

src/xtgeo/surface/_regsurf_import.py

@@ -3,12 +3,12 @@
 from __future__ import annotations
 
 import json
-import mmap
 from struct import unpack
 from typing import TYPE_CHECKING
 
 import h5py
 import numpy as np
+import surfio
 
 import xtgeo.common.sys as xsys
 from xtgeo import _cxtgeo
@@ -50,31 +50,20 @@ def import_irap_binary(mfile: FileWrapper, values: bool = True, **_):
     if mfile.memstream:
         mfile.file.seek(0)
         buf = mfile.file.read()
+        surf = surfio.IrapSurface.from_binary_buffer(buf)
     else:
-        with open(mfile.file, "rb") as fhandle:
-            buf = mmap.mmap(fhandle.fileno(), 0, access=mmap.ACCESS_READ)
-
-    # Ensure buffer is large enough for header
-    header_size = 100
-    if len(buf) < header_size:
-        raise ValueError("Buffer size is too small for header")
-
-    # unpack header with big-endian format string (cf. docstring info)
-    hed = np.frombuffer(
-        buf[:header_size],
-        dtype=">i4,>i4,>i4,>f4,>f4,>f4,>f4,>f4,>f4,>i4,"  # <32> IDFLAG NY XORI ... <32>
-        + ">i4,>i4,>f4,>f4,>f4,>i4,"  # <16> NX ROT X0ORI Y0ORI<16>
-        + ">i4,>i4,>i4,>i4,>i4,>i4,>i4,>i4,>i4",  # <28> 0 0 0 0 0 0 0 <28>
-    )
+        surf = surfio.IrapSurface.from_binary_file(mfile.file)
+
+    header = surf.header
 
     args = {}
-    args["nrow"] = int(hed[0][2])
-    args["xori"] = float(hed[0][3])
-    args["yori"] = float(hed[0][5])
-    args["xinc"] = float(hed[0][7])
-    args["yinc"] = float(hed[0][8])
-    args["ncol"] = int(hed[0][11])
-    args["rotation"] = float(hed[0][12])
+    args["nrow"] = header.nrow
+    args["xori"] = header.xori
+    args["yori"] = header.yori
+    args["xinc"] = header.xinc
+    args["yinc"] = header.yinc
+    args["ncol"] = header.ncol
+    args["rotation"] = header.rot
 
     args["yflip"] = 1
     if args["yinc"] < 0.0:
@@ -84,32 +73,15 @@ def import_irap_binary(mfile: FileWrapper, values: bool = True, **_):
     if not values:
         return args
 
-    # Values: traverse through data blocks
-    stv = header_size  # Starting byte
-    datav = []
-
-    while stv < len(buf):
-        # start block integer - number of bytes of floats in following block
-        blockv = np.frombuffer(buf[stv : stv + 4], dtype=">i4")[0]
-        stv += 4
-        # floats
-        datav.append(np.frombuffer(buf[stv : stv + blockv], dtype=">f4"))
-        stv += blockv
-        # end block integer not needed really
-        stv += 4
-
-    values = np.hstack(datav)
-    values = np.reshape(values, (args["ncol"], args["nrow"]), order="F")
-    values = np.array(values, order="C")
+    values = surf.values
     values = np.ma.masked_greater_equal(values, UNDEF_MAP_IRAPB)
     args["values"] = np.ma.masked_invalid(values)
 
-    del buf
     return args
 
 
 def import_irap_ascii(mfile: FileWrapper, **_):
-    """Import Irap in pure python code, suitable for memstreams, and now efficient.
+    """Import Irap ascii which has the following format:
     -996  2010      5.000000      5.000000
     461587.553724   467902.553724  5927061.430176  5937106.430176
     1264       30.000011   461587.553724  5927061.430176
@@ -124,33 +96,30 @@ def import_irap_ascii(mfile: FileWrapper, **_):
     if mfile.memstream:
         mfile.file.seek(0)
         buf = mfile.file.read().decode()
+        surface = surfio.IrapSurface.from_ascii_string(buf)
+        del buf
     else:
-        with open(mfile.file) as fhandle:
-            buf = fhandle.read()
-
-    buf = buf.split(maxsplit=19)
-    args = {}
-    args["nrow"] = int(buf[1])
-    args["xinc"] = float(buf[2])
-    args["yinc"] = float(buf[3])
-    args["xori"] = float(buf[4])
-    args["yori"] = float(buf[6])
-    args["ncol"] = int(buf[8])
-    args["rotation"] = float(buf[9])
-
-    nvalues = args["nrow"] * args["ncol"]
-    values = np.fromstring(buf[19], dtype=np.double, count=nvalues, sep=" ")
-
-    values = np.reshape(values, (args["ncol"], args["nrow"]), order="F")
-    values = np.array(values, order="C")
-    args["values"] = np.ma.masked_greater_equal(values, UNDEF_MAP_IRAPA)
+        surface = surfio.IrapSurface.from_ascii_file(str(mfile.file))
+
+    values = surface.values
+    values = np.ma.masked_greater_equal(values, UNDEF_MAP_IRAPA)
+    values = np.ma.masked_invalid(values)
+    args = {
+        "nrow": surface.header.nrow,
+        "xinc": surface.header.xinc,
+        "yinc": surface.header.yinc,
+        "xori": surface.header.xori,
+        "yori": surface.header.yori,
+        "ncol": surface.header.ncol,
+        "rotation": surface.header.rot,
+        "values": surface.values,
+    }
 
     args["yflip"] = 1
     if args["yinc"] < 0.0:
         args["yinc"] *= -1
         args["yflip"] = -1
 
-    del buf
     return args