Skip to content

Commit

Permalink
Integrate NumPy
Browse files Browse the repository at this point in the history
  • Loading branch information
@gerlero
gerlero committed Dec 6, 2024
1 parent ecbfb07 commit 1430e49
Show file tree
Hide file tree
Showing 11 changed files with 213 additions and 243 deletions.
64 changes: 25 additions & 39 deletions foamlib/_files/_files.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
from __future__ import annotations

import os
import sys
from copy import deepcopy
from typing import Any, Optional, Tuple, Union, cast
Expand All @@ -15,6 +14,8 @@
else:
from typing import Iterator, Mapping, MutableMapping, Sequence

import numpy as np

from ._io import FoamFileIO
from ._serialization import Kind, dumps, normalize
from ._types import (
Expand All @@ -27,7 +28,6 @@
File,
MutableEntry,
)
from ._util import is_sequence


class FoamFile(
Expand Down Expand Up @@ -228,50 +228,36 @@ def __setitem__(self, keywords: str | tuple[str, ...] | None, data: Entry) -> No
or keywords[2].endswith("Gradient")
)
):
if self.format == "binary":
arch = self.get(("FoamFile", "arch"), default=None)
assert arch is None or isinstance(arch, str)
if (arch is not None and "scalar=32" in arch) or (
arch is None
and os.environ.get("WM_PRECISION_OPTION", default="DP") == "SP"
):
kind = Kind.SINGLE_PRECISION_BINARY_FIELD
else:
kind = Kind.DOUBLE_PRECISION_BINARY_FIELD
else:
kind = Kind.ASCII_FIELD
kind = (
Kind.BINARY_FIELD if self.format == "binary" else Kind.ASCII_FIELD
)
elif keywords == ("dimensions",):
kind = Kind.DIMENSIONS

if (
kind
in (
Kind.ASCII_FIELD,
Kind.DOUBLE_PRECISION_BINARY_FIELD,
Kind.SINGLE_PRECISION_BINARY_FIELD,
)
kind in (Kind.ASCII_FIELD, Kind.BINARY_FIELD)
) and self.class_ == "dictionary":
if isinstance(data, (int, float)):
self.class_ = "volScalarField"

elif is_sequence(data) and data:
if isinstance(data[0], (int, float)):
if len(data) == 3:
self.class_ = "volVectorField"
elif len(data) == 6:
self.class_ = "volSymmTensorField"
elif len(data) == 9:
self.class_ = "volTensorField"
elif (
is_sequence(data[0])
and data[0]
and isinstance(data[0][0], (int, float))
):
if len(data[0]) == 3:
try:
shape = np.shape(data) # type: ignore [arg-type]
except ValueError:
pass
else:
if not shape:
self.class_ = "volScalarField"
elif shape == (3,):
self.class_ = "volVectorField"
elif shape == (6,):
self.class_ = "volSymmTensorField"
elif shape == (9,):
self.class_ = "volTensorField"
elif len(shape) == 1:
self.class_ = "volScalarField"
elif len(shape) == 2:
if shape[1] == 3:
self.class_ = "volVectorField"
elif len(data[0]) == 6:
elif shape[1] == 6:
self.class_ = "volSymmTensorField"
elif len(data[0]) == 9:
elif shape[1] == 9:
self.class_ = "volTensorField"

parsed = self._get_parsed(missing_ok=True)
Expand Down
61 changes: 24 additions & 37 deletions foamlib/_files/_parsing.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
from __future__ import annotations

import array
import re
import sys
from enum import Enum, auto
Expand All @@ -16,6 +15,7 @@
else:
EllipsisType = type(...)

import numpy as np
from pyparsing import (
Combine,
Dict,
Expand Down Expand Up @@ -47,13 +47,16 @@ class _Tensor(Enum):
TENSOR = auto()

@property
def shape(self) -> tuple[int, ...]:
return {
_Tensor.SCALAR: (),
_Tensor.VECTOR: (3,),
_Tensor.SYMM_TENSOR: (6,),
_Tensor.TENSOR: (9,),
}[self]
def shape(self) -> tuple[()] | tuple[int]:
if self == _Tensor.SCALAR:
return ()
if self == _Tensor.VECTOR:
return (3,)
if self == _Tensor.SYMM_TENSOR:
return (6,)
if self == _Tensor.TENSOR:
return (9,)
raise NotImplementedError

@property
def size(self) -> int:
Expand Down Expand Up @@ -84,7 +87,7 @@ def parser(self) -> ParserElement:
Literal("(").suppress()
+ Group(common.ieee_float[self.size], aslist=True)
+ Literal(")").suppress()
)
).add_parse_action(lambda tks: np.array(tks[0], dtype=float))

def __str__(self) -> str:
return {
Expand Down Expand Up @@ -116,40 +119,22 @@ def _list_of(entry: ParserElement) -> ParserElement:

def _parse_ascii_field(
s: str, tensor_kind: _Tensor, *, ignore: Regex | None
) -> list[float] | list[list[float]]:
values = [
float(v)
for v in (re.sub(ignore.re, " ", s) if ignore is not None else s)
.replace("(", " ")
.replace(")", " ")
.split()
]

if tensor_kind == _Tensor.SCALAR:
return values
) -> np.ndarray[tuple[int] | tuple[int, int], np.dtype[np.float64]]:
if ignore is not None:
s = re.sub(ignore.re, " ", s)
s = s.replace("(", " ").replace(")", " ")

return [
values[i : i + tensor_kind.size]
for i in range(0, len(values), tensor_kind.size)
]
return np.fromstring(s, dtype=float, sep=" ").reshape(-1, *tensor_kind.shape)


def _unpack_binary_field(
b: bytes, tensor_kind: _Tensor, *, length: int
) -> list[float] | list[list[float]]:
) -> np.ndarray[tuple[int] | tuple[int, int], np.dtype[np.float64 | np.float32]]:
float_size = len(b) / tensor_kind.size / length
assert float_size in (4, 8)

arr = array.array("f" if float_size == 4 else "d", b)
values = arr.tolist()

if tensor_kind == _Tensor.SCALAR:
return values

return [
values[i : i + tensor_kind.size]
for i in range(0, len(values), tensor_kind.size)
]
dtype = np.float32 if float_size == 4 else float
return np.frombuffer(b, dtype=dtype).reshape(-1, *tensor_kind.shape)


def _tensor_list(
Expand Down Expand Up @@ -187,7 +172,7 @@ def count_parse_action(tks: ParseResults) -> None:
)
| Regex(
rf"\((?s:.{{{length * tensor_kind.size * 8}}}|.{{{length * tensor_kind.size * 4}}})\)"
).set_parse_action(
).add_parse_action(
lambda tks: [
_unpack_binary_field(
tks[0][1:-1].encode("latin-1"), tensor_kind, length=length
Expand All @@ -196,7 +181,9 @@ def count_parse_action(tks: ParseResults) -> None:
)
| (
Literal("{").suppress() + tensor_kind.parser() + Literal("}").suppress()
).set_parse_action(lambda tks: [[tks[0]] * length])
).add_parse_action(
lambda tks: [np.full((length, *tensor_kind.shape), tks[0], dtype=float)]
)
)

count = common.integer.copy().add_parse_action(count_parse_action)
Expand Down
114 changes: 48 additions & 66 deletions foamlib/_files/_serialization.py
View file
Original file line number Diff line number Diff line change
@@ -1,34 +1,25 @@
from __future__ import annotations

import array
import itertools
import sys
from enum import Enum, auto
from typing import cast, overload
from typing import overload

if sys.version_info >= (3, 9):
from collections.abc import Mapping, Sequence
from collections.abc import Mapping
else:
from typing import Mapping, Sequence
from typing import Mapping

from ._parsing import parse_data
from ._types import Data, Dimensioned, DimensionSet, Entry
from ._util import is_sequence

try:
import numpy as np
import numpy as np

numpy = True
except ModuleNotFoundError:
numpy = False
from ._parsing import parse_data
from ._types import Data, Dimensioned, DimensionSet, Entry, is_sequence


class Kind(Enum):
DEFAULT = auto()
SINGLE_ENTRY = auto()
ASCII_FIELD = auto()
DOUBLE_PRECISION_BINARY_FIELD = auto()
SINGLE_PRECISION_BINARY_FIELD = auto()
BINARY_FIELD = auto()
DIMENSIONS = auto()


Expand All @@ -41,9 +32,29 @@ def normalize(data: Entry, *, kind: Kind = Kind.DEFAULT) -> Entry: ...


def normalize(data: Entry, *, kind: Kind = Kind.DEFAULT) -> Entry:
if numpy and isinstance(data, np.ndarray):
if kind in (Kind.ASCII_FIELD, Kind.BINARY_FIELD):
if is_sequence(data):
try:
arr = np.asarray(data)
except ValueError:
pass
else:
if not np.issubdtype(arr.dtype, np.floating):
arr = arr.astype(float)

if arr.ndim == 1 or (arr.ndim == 2 and arr.shape[1] in (3, 6, 9)):
return arr

return data

if isinstance(data, int):
return float(data)

return data

if isinstance(data, np.ndarray):
ret = data.tolist()
assert isinstance(ret, list)
assert isinstance(ret, (int, float, list))
return ret

if isinstance(data, Mapping):
Expand All @@ -55,7 +66,6 @@ def normalize(data: Entry, *, kind: Kind = Kind.DEFAULT) -> Entry:
and len(data) <= 7
and all(isinstance(d, (int, float)) for d in data)
):
data = cast(Sequence[float], data)
return DimensionSet(*data)

if isinstance(data, tuple) and kind == Kind.SINGLE_ENTRY and len(data) == 2:
Expand All @@ -65,17 +75,12 @@ def normalize(data: Entry, *, kind: Kind = Kind.DEFAULT) -> Entry:
if is_sequence(data) and (kind == Kind.SINGLE_ENTRY or not isinstance(data, tuple)):
return [normalize(d, kind=Kind.SINGLE_ENTRY) for d in data]

if isinstance(data, Dimensioned):
value = normalize(data.value, kind=Kind.SINGLE_ENTRY)
assert isinstance(value, (int, float, list))
return Dimensioned(value, data.dimensions, data.name)

if isinstance(data, str):
return parse_data(data)

if isinstance(
data,
(int, float, bool, tuple, DimensionSet),
(int, float, bool, tuple, DimensionSet, Dimensioned),
):
return data

Expand Down Expand Up @@ -107,58 +112,35 @@ def dumps(
if isinstance(data, DimensionSet):
return b"[" + b" ".join(dumps(v) for v in data) + b"]"

if kind in (
Kind.ASCII_FIELD,
Kind.DOUBLE_PRECISION_BINARY_FIELD,
Kind.SINGLE_PRECISION_BINARY_FIELD,
) and (
isinstance(data, (int, float))
or (
is_sequence(data)
and data
and isinstance(data[0], (int, float))
and len(data) in (3, 6, 9)
)
if kind in (Kind.ASCII_FIELD, Kind.BINARY_FIELD) and (
isinstance(data, (int, float, np.ndarray))
):
return b"uniform " + dumps(data, kind=Kind.SINGLE_ENTRY)

if kind in (
Kind.ASCII_FIELD,
Kind.DOUBLE_PRECISION_BINARY_FIELD,
Kind.SINGLE_PRECISION_BINARY_FIELD,
) and is_sequence(data):
if data and isinstance(data[0], (int, float)):
shape = np.shape(data)
if shape in ((), (3,), (6,), (9,)):
return b"uniform " + dumps(data, kind=Kind.SINGLE_ENTRY)

assert isinstance(data, np.ndarray)
ndim = len(shape)
if ndim == 1:
tensor_kind = b"scalar"
elif is_sequence(data[0]) and data[0] and isinstance(data[0][0], (int, float)):
if len(data[0]) == 3:

elif ndim == 2:
if shape[1] == 3:
tensor_kind = b"vector"
elif len(data[0]) == 6:
elif shape[1] == 6:
tensor_kind = b"symmTensor"
elif len(data[0]) == 9:
elif shape[1] == 9:
tensor_kind = b"tensor"
else:
return dumps(data)

else:
return dumps(data)

if kind in (
Kind.DOUBLE_PRECISION_BINARY_FIELD,
Kind.SINGLE_PRECISION_BINARY_FIELD,
):
typecode = "f" if kind == Kind.SINGLE_PRECISION_BINARY_FIELD else "d"
if tensor_kind == b"scalar":
data = cast(Sequence[float], data)
contents = b"(" + array.array(typecode, data).tobytes() + b")"
else:
data = cast(Sequence[Sequence[float]], data)
contents = (
b"("
+ array.array(
typecode, itertools.chain.from_iterable(data)
).tobytes()
+ b")"
)
if kind == Kind.BINARY_FIELD:
contents = b"(" + data.tobytes() + b")"
else:
assert kind == Kind.ASCII_FIELD
contents = dumps(data, kind=Kind.SINGLE_ENTRY)

return b"nonuniform List<" + tensor_kind + b"> " + dumps(len(data)) + contents
Expand Down
Loading

0 comments on commit 1430e49

Please sign in to comment.