From a777a29922a5e61088d05778a57c2282a5a70285 Mon Sep 17 00:00:00 2001 From: JB Lovland Date: Thu, 30 Nov 2023 17:54:14 +0100 Subject: [PATCH] CLN: Add types gridprop-import-xtgcpgeom --- mypy.ini | 3 - pyproject.toml | 5 +- src/xtgeo/grid3d/_grid_import_xtgcpgeom.py | 175 ++++++++++++++++----- 3 files changed, 135 insertions(+), 48 deletions(-) diff --git a/mypy.ini b/mypy.ini index 9fdcd5caa..3f5779acd 100644 --- a/mypy.ini +++ b/mypy.ini @@ -69,9 +69,6 @@ ignore_errors = True [mypy-xtgeo.grid3d._grid_import_roff] ignore_errors = True -[mypy-xtgeo.grid3d._grid_import_xtgcpgeom] -ignore_errors = True - [mypy-xtgeo.grid3d._grid_refine] ignore_errors = True diff --git a/pyproject.toml b/pyproject.toml index 87989d739..7ce274951 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,19 +48,20 @@ classifiers = [ dynamic = ["version"] dependencies = [ "deprecation", - "resfo>=4.0.0", "h5py>=3", "hdf5plugin>=2.3", "joblib", "matplotlib>=3.3", "numpy<2", "pandas>=1.1", + "resfo>=4.0.0", "roffio>=0.0.2", "scipy>=1.5", "segyio>1.8.0", "shapely>=1.6.2", - "tables<3.9;platform_system != 'Darwin' and python_version == '3.8'", # TODO: mac... "tables;platform_system != 'Darwin' and python_version > '3.8'", # TODO: mac... + "tables<3.9;platform_system != 'Darwin' and python_version == '3.8'", # TODO: mac... + "typing_extensions", ] [project.urls] diff --git a/src/xtgeo/grid3d/_grid_import_xtgcpgeom.py b/src/xtgeo/grid3d/_grid_import_xtgcpgeom.py index 793f0ceb5..aa60cc229 100644 --- a/src/xtgeo/grid3d/_grid_import_xtgcpgeom.py +++ b/src/xtgeo/grid3d/_grid_import_xtgcpgeom.py @@ -1,19 +1,68 @@ -# coding: utf-8 """Private module, Grid Import private functions for xtgeo based formats.""" +from __future__ import annotations import json from struct import unpack +from typing import TYPE_CHECKING, Literal, Tuple, TypedDict, Union, overload import h5py import numpy as np +from typing_extensions import NotRequired, TypeAlias import xtgeo.common.sys as xsys from xtgeo.common import null_logger +from xtgeo.grid3d._gridprop_import_xtgcpprop import _read_filelike logger = null_logger(__name__) +if TYPE_CHECKING: + from xtgeo.common.sys import _XTGeoFile -def convert_subgrids(sdict): +LMIN: TypeAlias = Literal["min"] +LMAX: TypeAlias = Literal["max"] +IJKRANGE: TypeAlias = Tuple[ + Union[int, LMIN], + Union[int, LMAX], + Union[int, LMIN], + Union[int, LMAX], + Union[int, LMIN], + Union[int, LMAX], +] + + +class _RESULT(TypedDict): + coordsv: np.ndarray + zcornsv: np.ndarray + actnumsv: np.ndarray + subgrids: NotRequired[dict[str, range]] + + +class __REQUIRED_(TypedDict): + xshift: np.ndarray + xscale: np.ndarray + yshift: np.ndarray + yscale: np.ndarray + zshift: np.ndarray + zscale: np.ndarray + subgrids: NotRequired[dict[str, int]] + + +class _META(TypedDict): + _required_: __REQUIRED_ + + +@overload +def convert_subgrids(sdict: dict[str, int]) -> dict[str, range]: + ... + + +@overload +def convert_subgrids(sdict: None) -> None: + ... + + +def convert_subgrids(sdict: dict[str, int] | None) -> dict[str, range] | None: + # flake8: noqa: F841 """Convert a simplified ordered dictionary to subgrid required by Grid. @@ -29,18 +78,17 @@ def convert_subgrids(sdict): if not isinstance(sdict, dict): raise ValueError("Input sdict is not an dict") - newsub = dict() - - inn1 = 1 - for name, nsub in sdict.items(): - inn2 = inn1 + nsub - newsub[name] = range(inn1, inn2) - inn1 = inn2 - - return newsub + start = 1 + return {name: range(start, start := start + nsub) for name, nsub in sdict.items()} -def handle_metadata(result, meta, ncol, nrow, nlay): +def handle_metadata( + result: _RESULT, + meta: _META, + ncol: int, + nrow: int, + nlay: int, +) -> None: # meta _optional_ *may* contain xshift, xscale etc which in case must be taken # into account coordsv = result["coordsv"] @@ -66,21 +114,22 @@ def handle_metadata(result, meta, ncol, nrow, nlay): result["coordsv"] = coordsv.reshape((nncol, nnrow, 6)).astype(np.float64) result["zcornsv"] = zcornsv.reshape((nncol, nnrow, nnlay, 4)).astype(np.float32) result["actnumsv"] = actnumsv.reshape((ncol, nrow, nlay)).astype(np.int32) + if "subgrids" in req: result["subgrids"] = convert_subgrids(req["subgrids"]) def import_xtgcpgeom( - mfile, mmap=False -): # pylint: disable=too-many-locals, too-many-statements + mfile: _XTGeoFile, + mmap: bool = False, +) -> _RESULT: # pylint: disable=too-many-locals, too-many-statements """Using pure python for experimental grid geometry import.""" # offset = 36 - with open(mfile.file, "rb") as fhandle: - buf = fhandle.read(offset) + with _read_filelike(mfile.file, size=offset) as f: + # unpack header + swap, magic, nformat, ncol, nrow, nlay = unpack("= i i i q q q", f) - # unpack header - swap, magic, nformat, ncol, nrow, nlay = unpack("= i i i q q q", buf) nncol = ncol + 1 nnrow = nrow + 1 nnlay = nlay + 1 @@ -95,41 +144,57 @@ def import_xtgcpgeom( coordfmt, zcornfmt, actnumfmt = [int(nbyte) for nbyte in str(nformat)] - dtype_coordsv = "float" + str(coordfmt * 8) - dtype_zcornsv = "float" + str(zcornfmt * 8) - dtype_actnumv = "int" + str(actnumfmt * 8) + dtype_coordsv = f"float{coordfmt * 8}" + dtype_zcornsv = f"float{zcornfmt * 8}" + dtype_actnumv = f"int{actnumfmt * 8}" ncoord = nncol * nnrow * 6 nzcorn = nncol * nnrow * nnlay * 4 nactnum = ncol * nrow * nlay - result = dict() + # read numpy arrays from file - result["coordsv"] = xsys.npfromfile( - mfile.file, dtype=dtype_coordsv, count=ncoord, offset=offset, mmap=mmap + coordsv = xsys.npfromfile( + mfile.file, + dtype=dtype_coordsv, + count=ncoord, + offset=offset, + mmap=mmap, ) + newoffset = offset + ncoord * coordfmt - result["zcornsv"] = xsys.npfromfile( - mfile.file, dtype=dtype_zcornsv, count=nzcorn, offset=newoffset, mmap=mmap + zcornsv = xsys.npfromfile( + mfile.file, + dtype=dtype_zcornsv, + count=nzcorn, + offset=newoffset, + mmap=mmap, ) + newoffset += nzcorn * zcornfmt - result["actnumsv"] = xsys.npfromfile( - mfile.file, dtype=dtype_actnumv, count=nactnum, offset=newoffset, mmap=mmap + actnumsv = xsys.npfromfile( + mfile.file, + dtype=dtype_actnumv, + count=nactnum, + offset=newoffset, + mmap=mmap, ) + newoffset += nactnum * actnumfmt + result = _RESULT(coordsv=coordsv, zcornsv=zcornsv, actnumsv=actnumsv) # read metadata which will be at position offet + nfloat*narr +13 - pos = newoffset + 13 - with open(mfile.file, "rb") as fhandle: - fhandle.seek(pos) - jmeta = fhandle.read().decode() - - meta = json.loads(jmeta, object_pairs_hook=dict) + with _read_filelike(mfile.file, seek=newoffset + 13) as f: + meta = json.loads(f, object_pairs_hook=dict) handle_metadata(result, meta, ncol, nrow, nlay) return result -def import_hdf5_cpgeom(mfile, ijkrange=None, zerobased=False): +def import_hdf5_cpgeom( + mfile: _XTGeoFile, + ijkrange: IJKRANGE | None = None, + zerobased: bool = False, +) -> _RESULT: """Experimental grid geometry import using hdf5.""" # with h5py.File(mfile.name, "r") as h5h: @@ -157,24 +222,36 @@ def import_hdf5_cpgeom(mfile, ijkrange=None, zerobased=False): "subgrids" in meta["_required_"] and meta["_required_"]["subgrids"] is not None ): + *_, k1, k2 = ijkrange meta["_required_"]["subgrids"] = filter_subgrids_partial( - meta["_required_"]["subgrids"], *ijkrange[-2:], nlay, zerobased + meta["_required_"]["subgrids"], + k1, + k2, + nlay, + zerobased, ) else: incoord = grp["coord"][:, :, :] inzcorn = grp["zcorn"][:, :, :, :] inactnum = grp["actnum"][:, :, :] - result = {} - result["coordsv"] = incoord.astype("float64") - result["zcornsv"] = inzcorn.astype("float32") - result["actnumsv"] = inactnum.astype("float32") + result = _RESULT( + coordsv=incoord.astype("float64"), + zcornsv=inzcorn.astype("float32"), + actnumsv=inactnum.astype("float32"), + ) handle_metadata(result, meta, ncol, nrow, nlay) return result -def filter_subgrids_partial(subgrids, k1, k2, nlay, zerobased): +def filter_subgrids_partial( + subgrids: dict[str, int], + k1: int | LMIN, + k2: int | LMAX, + nlay: int, + zerobased: bool, +) -> dict[str, int]: """ Filters and truncates the subgrids of the global grid so that they refer to the filtered grid. @@ -212,7 +289,7 @@ def filter_subgrids_partial(subgrids, k1, k2, nlay, zerobased): k1 -= 1 k2 -= 1 - partial_subgrid = dict() + partial_subgrid: dict[str, int] = dict() start = 0 for key, value in subgrids.items(): end = value + start @@ -224,7 +301,19 @@ def filter_subgrids_partial(subgrids, k1, k2, nlay, zerobased): return partial_subgrid -def _partial_read(h5h, req, ijkrange, zerobased): +def _partial_read( + h5h: dict[str, np.ndarray], + req: dict[str, int], + ijkrange: tuple[ + int | LMIN, + int | LMAX, + int | LMIN, + int | LMAX, + int | LMIN, + int | LMAX, + ], + zerobased: bool, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, int, int, int]: """Read a partial IJ range.""" ncol = req["ncol"] nrow = req["nrow"]