implement/test array parsing

flopy4/mf6/codec/reader/__init__.py

@@ -2,8 +2,8 @@
 from pathlib import Path
 from typing import Any
 
-from flopy4.mf6.codec.reader.parser import make_generic_parser
-from flopy4.mf6.codec.reader.transformer import GenericTransformer
+from flopy4.mf6.codec.reader.parser import make_basic_parser
+from flopy4.mf6.codec.reader.transformer import BasicTransformer
 
 
 def load(path: str | PathLike) -> Any:
@@ -41,6 +41,6 @@ def loads(data: str) -> Any:
         Parsed MF6 input file structure
     """
 
-    parser = make_generic_parser()
-    transformer = GenericTransformer()
+    parser = make_basic_parser()
+    transformer = BasicTransformer()
     return transformer.transform(parser.parse(data))

flopy4/mf6/codec/reader/grammar/array.lark

@@ -0,0 +1,34 @@
+start: readarray
+readarray: (single_array | layered_array)
+single_array: [netcdf] array
+layered_array: layered [netcdf] array+
+layered: "layered"i
+netcdf: "netcdf"i
+array: control [data]
+control: constant | internal | external
+constant: "constant"i _number
+internal: "internal"i [factor] [iprn]
+external: "open/close"i filename [factor] [binary] [iprn]
+factor: "factor"i _number
+iprn: "iprn"i _signed_integer
+binary: "(" "binary"i ")"
+filename: ESCAPED_STRING | _word
+data: _number+
+_number: SIGNED_NUMBER | NUMBER
+_signed_integer: SIGNED_INT | INT
+_word: /[a-zA-Z0-9._'~,-\\(\\)]+/
+
+%import common.NEWLINE -> _NL
+%import common.WS
+%import common.WS_INLINE
+%import common.CNAME
+%import common.WORD
+%import common.NUMBER
+%import common.INT
+%import common.SH_COMMENT
+%import common.ESCAPED_STRING
+%import common.SIGNED_NUMBER
+%import common.SIGNED_INT
+
+%ignore WS
+%ignore SH_COMMENT

flopy4/mf6/codec/reader/grammar/mf6.lark → flopy4/mf6/codec/reader/grammar/basic.lark

@@ -15,5 +15,6 @@ word: /[a-zA-Z0-9._'~,-\\(\\)]+/
 %import common.INT
 %import common.SH_COMMENT
 %import common._STRING_INNER
+
 %ignore WS_INLINE
 %ignore SH_COMMENT

flopy4/mf6/codec/reader/parser.py

@@ -3,8 +3,15 @@
 from lark import Lark
 
 
-def make_generic_parser() -> Lark:
-    grammar_path = Path(__file__).parent / "grammar" / "mf6.lark"
+def make_basic_parser() -> Lark:
+    grammar_path = Path(__file__).parent / "grammar" / "basic.lark"
+    with open(grammar_path, "r") as f:
+        grammar = f.read()
+    return Lark(grammar, parser="lalr", debug=True)
+
+
+def make_array_parser() -> Lark:
+    grammar_path = Path(__file__).parent / "grammar" / "array.lark"
     with open(grammar_path, "r") as f:
         grammar = f.read()
     return Lark(grammar, parser="lalr", debug=True)

flopy4/mf6/codec/reader/transformer.py

@@ -1,12 +1,15 @@
+from pathlib import Path
 from typing import Any
 
+import numpy as np
+import xarray as xr
 from lark import Token, Transformer
 
 
-class GenericTransformer(Transformer):
+class BasicTransformer(Transformer):
     """
-    Generic transformer for MF6 input files. Works only with the generic
-    grammar. Returns structures of blocks consisting of lines of tokens.
+    Basic transformer for MF6 input files. Works only with the basic
+    grammar. Yields blocks simply as collections of lines of tokens.
     """
 
     def start(self, items: list[Any]) -> dict[str, Any]:
@@ -48,3 +51,123 @@ def CNAME(self, token: Token) -> str:
 
     def INT(self, token: Token) -> int:
         return int(token)
+
+
+class ArrayTransformer(Transformer):
+    """
+    Transformer for MF6 array input format. Returns xarray DataArrays
+    for internal/constant arrays and Path objects for external arrays,
+    inside a dictionary which also contains control information. This
+    is a first step towards a smarter parser/transformer for the full
+    MF6 input format specification.
+    """
+
+    def start(self, items: list[Any]) -> dict:
+        return items[0]
+
+    def readarray(self, items: list[Any]) -> dict:
+        infos = items[0]
+        if isinstance(infos, list):
+            data = xr.concat([info["data"] for info in infos if "data" in info], dim="layer")
+            return {
+                "control": [info["control"] for info in infos if "control" in info],
+                "data": data,
+                "attrs": {k: v for k, v in infos[0].items() if k not in ["data"]},
+                "dims": {"layer": len(infos)},
+            }
+        return infos
+
+    def single_array(self, items: list[Any]) -> dict:
+        netcdf = items[0]
+        info = items[-1]
+        if netcdf:
+            info["netcdf"] = netcdf
+        return ArrayTransformer.try_create_dataarray(info)
+
+    def layered_array(self, items: list[Any]) -> list[dict]:
+        netcdf = items[0]
+        infos = []
+        for info in items[2:]:
+            if info is None:
+                continue
+            if netcdf:
+                info["netcdf"] = netcdf
+            infos.append(ArrayTransformer.try_create_dataarray(info))
+        return infos
+
+    def array(self, items: list[Any]) -> dict[str, Any]:
+        control = items[0]
+        data = items[1] if len(items) > 1 else None
+        if (value := control.get("value", None)) is not None:
+            data = value
+        return {"control": control, "data": data}
+
+    def control(self, items: list[Any]) -> dict[str, Any]:
+        return items[0]
+
+    def constant(self, items: list[Any]) -> dict[str, Any]:
+        return {"type": "constant", "value": items[0]}
+
+    def internal(self, items: list[Any]) -> dict[str, Any]:
+        result = {"type": "internal"}
+        for item in items:
+            if item is not None:
+                result.update(item)
+        return result
+
+    def external(self, items: list[Any]) -> dict[str, Any]:
+        return {"type": "external", "value": items[0]}
+
+    def factor(self, items: list[Any]) -> dict[str, float]:
+        return {"factor": items[0]}
+
+    def iprn(self, items: list[Any]) -> dict[str, int]:
+        return {"iprn": items[0]}
+
+    def binary(self, items: list[Any]) -> dict[str, bool]:
+        return {"binary": True}
+
+    def filename(self, items: list[Any]) -> Path:
+        return Path(items[0])
+
+    def data(self, items: list[Any]) -> np.ndarray:
+        return np.array(items)
+
+    def netcdf(self, items: list[Any]) -> dict[str, bool]:
+        return {"netcdf": True}
+
+    def NUMBER(self, token: Token) -> int | float:
+        return float(token)
+
+    def SIGNED_NUMBER(self, token: Token) -> int | float:
+        return self.NUMBER(token)
+
+    def INT(self, token: Token) -> int:
+        return int(token)
+
+    def SIGNED_INT(self, token: Token) -> int:
+        return int(token)
+
+    def ESCAPED_STRING(self, token: Token) -> str:
+        # Remove quotes from escaped string
+        value = str(token)
+        if value.startswith('"') and value.endswith('"'):
+            return value[1:-1]
+        return value
+
+    @staticmethod
+    def try_create_dataarray(array_info: dict) -> dict:
+        """Create an xarray DataArray from MF6 array information."""
+        control = array_info["control"]
+        match control["type"]:
+            case "constant":
+                data = control["value"]
+                array_info["data"] = xr.DataArray(data=data)
+            case "internal":
+                data = array_info["data"]
+                factor = control.get("factor", 1.0)
+                data = data * factor
+                array_info["data"] = xr.DataArray(data=data)
+            case "external":
+                pass
+        return array_info