fullplot 0.0.1__py3-none-any.whl

fullplot/maps.py

@@ -0,0 +1,708 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable
+
+import datetime as _datetime
+import json
+
+import h5py
+import numpy as np
+
+from fullplot.hdf5 import dataset_names, hdf5_filename, safe_group_name
+
+
+_STRING_DTYPE = h5py.string_dtype(encoding="utf-8")
+
+
+class FullPlotMapError(Exception):
+    """Base class for FullPlot map-generation errors."""
+
+
+class MapGenerationError(FullPlotMapError, RuntimeError):
+    """A generated map failed during evaluation, output discovery, or resume."""
+
+
+class MapOutputError(FullPlotMapError, TypeError):
+    """A map evaluate() function returned invalid outputs."""
+_RESERVED_GROUP_NAMES = {"axes", "outputs", "status"}
+
+
+@dataclass(frozen=True)
+class Axis:
+    """One independent variable for a generated HDF5 map.
+
+    An ``Axis`` defines one input direction of a tabulated map generated by
+    :func:`generate_map`. The axis ``name`` is important: it becomes the
+    keyword passed to the map-generation ``evaluate`` function and the input
+    name that must later be supplied to ``Map.from_hdf5``.
+
+    For example, an axis named ``"pressure"`` means that ``generate_map``
+    will call ``evaluate(pressure=...)`` and that the map is later loaded with
+    ``inputs={"pressure": some_state}``.
+
+    Parameters
+    ----------
+    name:
+        Public input name for this axis. It must be a non-empty string and
+        cannot contain ``"/"`` because it is also used as an HDF5 dataset name.
+    values:
+        Strictly increasing one-dimensional grid values in physical units.
+    units:
+        Optional units string stored in the HDF5 file for user inspection.
+        FullPlot does not perform unit conversion.
+    spacing:
+        How interpolation should treat this axis. Supported values are
+        ``"linear"``, ``"log"``, and ``"values"``. ``"log"`` stores physical
+        values but tells ``Map`` to interpolate versus ``log(value)``.
+
+    Notes
+    -----
+    Use an ``Axis`` for a quantity that should vary across the map. Use
+    ``constants`` in :func:`generate_map` for quantities that should be passed
+    to every ``evaluate`` call but are fixed for the generated map.
+    """
+
+    name: str
+    values: np.ndarray
+    units: str = ""
+    spacing: str = "values"
+
+    @classmethod
+    def linear(
+        cls,
+        name: str,
+        start: float,
+        stop: float,
+        count: int,
+        units: str = "",
+    ) -> "Axis":
+        """Create an evenly spaced axis using ``np.linspace``.
+
+        The returned axis includes both ``start`` and ``stop`` and has
+        ``spacing="linear"``. Use this for inputs that vary smoothly with
+        the physical value, such as temperature, mixture ratio, or area ratio.
+        """
+        return cls(
+            name=name,
+            values=np.linspace(start, stop, count, dtype=float),
+            units=units,
+            spacing="linear",
+        )
+
+    @classmethod
+    def log(
+        cls,
+        name: str,
+        start: float,
+        stop: float,
+        count: int,
+        units: str = "",
+    ) -> "Axis":
+        """Create a logarithmically spaced axis using ``np.geomspace``.
+
+        The stored axis values are still physical values, not their logarithms.
+        During lookup, ``Map`` applies ``log`` to both the stored axis values
+        and the runtime input value before interpolation. ``start`` and
+        ``stop`` must be positive. Use this for pressure or other quantities
+        spanning several orders of magnitude.
+        """
+        return cls(
+            name=name,
+            values=np.geomspace(start, stop, count, dtype=float),
+            units=units,
+            spacing="log",
+        )
+
+    @classmethod
+    def values(
+        cls,
+        name: str,
+        values,
+        units: str = "",
+    ) -> "Axis":
+        """Create an axis from explicit grid values.
+
+        Use this when the grid should follow measured points, hand-picked
+        breakpoints, map speed lines, or any other non-uniform spacing. The
+        values must be finite, one-dimensional, and strictly increasing after
+        validation.
+        """
+        return cls(
+            name=name,
+            values=np.asarray(values, dtype=float),
+            units=units,
+            spacing="values",
+        )
+
+
+
+def _validate_hdf5_name(name: str, label: str) -> None:
+    if not isinstance(name, str) or not name.strip():
+        raise ValueError(f"{label} names must be non-empty strings.")
+
+    if "/" in name:
+        raise ValueError(f"{label} name '{name}' cannot contain '/'.")
+
+
+def _validate_axis(axis: Axis) -> None:
+    _validate_hdf5_name(axis.name, "Axis")
+
+    values = np.asarray(axis.values, dtype=float)
+
+    if axis.spacing not in {"linear", "log", "values"}:
+        raise ValueError(f"Axis '{axis.name}' has unsupported spacing '{axis.spacing}'.")
+
+    if values.ndim != 1:
+        raise ValueError(f"Axis '{axis.name}' must be one-dimensional.")
+
+    if len(values) < 2:
+        raise ValueError(f"Axis '{axis.name}' requires at least two points.")
+
+    if not np.all(np.isfinite(values)):
+        raise ValueError(f"Axis '{axis.name}' values must all be finite.")
+
+    if np.any(np.diff(values) <= 0.0):
+        raise ValueError(f"Axis '{axis.name}' must be strictly increasing.")
+
+    if axis.spacing == "log" and np.any(values <= 0.0):
+        raise ValueError(f"Log axis '{axis.name}' must contain only positive values.")
+
+
+def _shape_from_axes(axes: list[Axis]) -> tuple[int, ...]:
+    return tuple(len(axis.values) for axis in axes)
+
+
+def _inputs_from_index(axes: list[Axis], index: tuple[int, ...], constants: dict) -> dict[str, float]:
+    inputs = dict(constants)
+
+    for axis, i in zip(axes, index):
+        inputs[axis.name] = float(axis.values[i])
+
+    return inputs
+
+
+def _json_default(value):
+    if isinstance(value, np.ndarray):
+        return value.tolist()
+
+    if isinstance(value, np.generic):
+        return value.item()
+
+    return repr(value)
+
+
+def _json_roundtrip(value):
+    return json.loads(json.dumps(value, default=_json_default))
+
+
+def _json_attr(attrs, name: str, default):
+    if name not in attrs:
+        return default
+
+    value = attrs[name]
+
+    if isinstance(value, bytes):
+        value = value.decode("utf-8")
+
+    if isinstance(value, str):
+        return json.loads(value)
+
+    return value
+
+
+
+def _validate_scalar_outputs(
+    values: dict,
+    outputs: list[str] | tuple[str, ...] | None = None,
+) -> dict[str, float]:
+    if not isinstance(values, dict):
+        raise MapOutputError("evaluate() must return a dictionary of scalar outputs.")
+
+    scalar_outputs: dict[str, float] = {}
+
+    for output_name, value in values.items():
+        try:
+            _validate_hdf5_name(output_name, "Output")
+        except ValueError as exc:
+            raise MapOutputError(str(exc)) from exc
+
+        if output_name in _RESERVED_GROUP_NAMES:
+            raise MapOutputError(f"evaluate() output name '{output_name}' is reserved.")
+
+        if isinstance(value, dict):
+            raise MapOutputError(
+                f"evaluate() output '{output_name}' must be a scalar number. "
+                "Nested dictionaries and vector outputs are not supported. "
+                "Return individual scalar entries instead."
+            )
+
+        if isinstance(value, (str, bytes, bool, np.bool_)):
+            raise MapOutputError(f"evaluate() output '{output_name}' must be numeric.")
+
+        array = np.asarray(value)
+
+        if array.shape != ():
+            raise MapOutputError(
+                f"evaluate() output '{output_name}' must be a scalar number. "
+                "Lists, tuples, arrays, and vector outputs are not supported."
+            )
+
+        try:
+            scalar_outputs[output_name] = float(array)
+        except Exception as exc:
+            raise MapOutputError(f"evaluate() output '{output_name}' must be numeric.") from exc
+
+    if outputs is not None:
+        missing = [
+            output_name
+            for output_name in outputs
+            if output_name not in scalar_outputs
+        ]
+
+        if missing:
+            raise MapOutputError(f"evaluate() did not return required outputs: {missing}")
+
+        scalar_outputs = {
+            output_name: scalar_outputs[output_name]
+            for output_name in outputs
+        }
+
+    return scalar_outputs
+
+
+def _write_axis_dataset(axes_group: h5py.Group, axis: Axis) -> None:
+    dataset = axes_group.create_dataset(
+        axis.name,
+        data=np.asarray(axis.values, dtype=float),
+    )
+
+    dataset.attrs["name"] = axis.name
+    dataset.attrs["units"] = axis.units
+    dataset.attrs["spacing"] = axis.spacing
+
+
+def _create_group(
+    file: h5py.File,
+    group: str,
+    axes: list[Axis],
+    output_names: list[str],
+    constants: dict,
+    metadata: dict,
+    fill_value: float,
+    compression: str | None,
+    compression_opts,
+) -> h5py.Group:
+    map_group = file.create_group(safe_group_name(group))
+
+    map_group.attrs["kind"] = "map"
+    map_group.attrs["name"] = group
+    map_group.attrs["map_format"] = "fullflow-map-v3"
+    map_group.attrs["created_utc"] = _datetime.datetime.now(_datetime.UTC).isoformat()
+    map_group.attrs["axis_order"] = json.dumps([axis.name for axis in axes])
+    map_group.attrs["output_names"] = json.dumps(output_names)
+    map_group.attrs["constants"] = json.dumps(constants, default=_json_default)
+    map_group.attrs["metadata"] = json.dumps(metadata, default=_json_default)
+
+    axes_group = map_group.create_group("axes")
+    outputs_group = map_group.create_group("outputs")
+    status_group = map_group.create_group("status")
+
+    for axis in axes:
+        _write_axis_dataset(axes_group, axis)
+
+    shape = _shape_from_axes(axes)
+
+    if compression is None:
+        compression_opts = None
+
+    for output_name in output_names:
+        outputs_group.create_dataset(
+            output_name,
+            shape=shape,
+            dtype="f8",
+            chunks=True,
+            compression=compression,
+            compression_opts=compression_opts,
+            fillvalue=fill_value,
+        )
+
+    status_group.create_dataset(
+        "success",
+        shape=shape,
+        dtype="?",
+        chunks=True,
+        compression=compression,
+        compression_opts=compression_opts,
+        fillvalue=False,
+    )
+
+    status_group.create_dataset(
+        "message",
+        shape=shape,
+        dtype=_STRING_DTYPE,
+        chunks=True,
+    )
+
+    return map_group
+
+
+def _check_existing_group(
+    map_group: h5py.Group,
+    axes: list[Axis],
+    outputs: list[str] | tuple[str, ...] | None,
+    constants: dict,
+    metadata: dict,
+) -> list[str]:
+    shape = _shape_from_axes(axes)
+
+    if "axes" not in map_group or not isinstance(map_group["axes"], h5py.Group):
+        raise MapGenerationError("Existing map group is not a fullflow-map-v3 group and cannot be resumed.")
+
+    if "outputs" not in map_group or not isinstance(map_group["outputs"], h5py.Group):
+        raise MapGenerationError("Existing map group is missing required group 'outputs'.")
+
+    if "status" not in map_group or not isinstance(map_group["status"], h5py.Group):
+        raise MapGenerationError("Existing map group is missing required group 'status'.")
+
+    axes_group = map_group["axes"]
+    outputs_group = map_group["outputs"]
+    status_group = map_group["status"]
+
+    for axis in axes:
+        if axis.name not in axes_group:
+            raise MapGenerationError(f"Existing map group is missing axis '{axis.name}'.")
+
+        axis_values = np.asarray(axes_group[axis.name][()], dtype=float)
+
+        if not np.array_equal(axis_values, np.asarray(axis.values, dtype=float)):
+            raise MapGenerationError(f"Existing map axis '{axis.name}' does not match requested axis.")
+
+    axis_order = json.loads(map_group.attrs.get("axis_order", "[]"))
+
+    if axis_order and axis_order != [axis.name for axis in axes]:
+        raise MapGenerationError("Existing map axis order does not match requested axis order.")
+
+    for axis in axes:
+        existing_spacing = axes_group[axis.name].attrs.get("spacing", "values")
+        if isinstance(existing_spacing, bytes):
+            existing_spacing = existing_spacing.decode("utf-8")
+        if str(existing_spacing).lower() != str(axis.spacing).lower():
+            raise MapGenerationError(
+                f"Existing map axis '{axis.name}' has spacing {existing_spacing!r}; "
+                f"requested spacing {axis.spacing!r}."
+            )
+
+    existing_constants = _json_attr(map_group.attrs, "constants", {})
+    requested_constants = _json_roundtrip(constants)
+    if existing_constants != requested_constants:
+        raise MapGenerationError(
+            "Existing map constants do not match requested constants. "
+            "Use overwrite=True to regenerate the map with different constants."
+        )
+
+    existing_metadata = _json_attr(map_group.attrs, "metadata", {})
+    requested_metadata = _json_roundtrip(metadata)
+    if existing_metadata != requested_metadata:
+        raise MapGenerationError(
+            "Existing map metadata does not match requested metadata. "
+            "Use overwrite=True to regenerate the map with different metadata."
+        )
+
+    output_names = list(outputs) if outputs is not None else dataset_names(outputs_group, set())
+
+    if not output_names:
+        raise MapGenerationError("Existing map group does not contain any output datasets.")
+
+    missing = [
+        output_name
+        for output_name in output_names
+        if output_name not in outputs_group
+    ]
+
+    if missing:
+        raise MapGenerationError(f"Existing map group is missing requested outputs: {missing}")
+
+    for output_name in output_names:
+        if outputs_group[output_name].shape != shape:
+            raise MapGenerationError(
+                f"Existing output '{output_name}' has shape {outputs_group[output_name].shape}; "
+                f"expected {shape}."
+            )
+
+    if "success" not in status_group:
+        raise MapGenerationError("Existing map group is missing required dataset 'status/success'.")
+
+    if "message" not in status_group:
+        raise MapGenerationError("Existing map group is missing required dataset 'status/message'.")
+
+    if status_group["success"].shape != shape:
+        raise MapGenerationError("Existing status/success shape does not match requested map axes.")
+
+    return output_names
+
+
+def _discover_outputs(
+    axes: list[Axis],
+    constants: dict,
+    evaluate: Callable,
+    outputs: list[str] | tuple[str, ...] | None,
+    raise_errors: bool,
+) -> tuple[list[str], tuple[int, ...], dict[str, float], list[tuple[tuple[int, ...], str]]]:
+    shape = _shape_from_axes(axes)
+    failures: list[tuple[tuple[int, ...], str]] = []
+
+    for index in np.ndindex(shape):
+        inputs = _inputs_from_index(axes, index, constants)
+
+        try:
+            values = _validate_scalar_outputs(evaluate(**inputs), outputs)
+            return list(values.keys()), index, values, failures
+        except MapOutputError:
+            raise
+        except Exception as exc:
+            if raise_errors:
+                raise
+
+            failures.append((index, f"{type(exc).__name__}: {exc}"))
+
+    raise MapGenerationError("Could not discover map outputs because every evaluated point failed.")
+
+
+def _write_failure(
+    map_group: h5py.Group,
+    index: tuple[int, ...],
+    output_names: list[str],
+    fill_value: float,
+    failure_message: str,
+) -> None:
+    outputs_group = map_group["outputs"]
+    status_group = map_group["status"]
+
+    for output_name in output_names:
+        outputs_group[output_name][index] = fill_value
+
+    status_group["success"][index] = False
+    status_group["message"][index] = failure_message
+
+
+def generate_map(
+    filename: str | Path,
+    axes: list[Axis] | tuple[Axis, ...],
+    evaluate: Callable,
+    group: str = "map",
+    outputs: list[str] | tuple[str, ...] | None = None,
+    constants: dict | None = None,
+    metadata: dict | None = None,
+    resume: bool = True,
+    overwrite: bool = False,
+    fill_value: float = np.nan,
+    compression: str | None = "gzip",
+    compression_opts=None,
+    flush_every: int = 25,
+    raise_errors: bool = False,
+) -> str:
+    """Generate a FullPlot HDF5 map compatible with FullFlow Map.from_hdf5().
+
+    The generated file can be loaded by FullFlow ``Map.from_hdf5`` or any reader that follows the same HDF5 map layout. ``generate_map`` evaluates a Python function on every
+    combination of axis values and writes one HDF5 dataset per scalar output.
+
+    Naming rules
+    ------------
+    Axis names are the input names. If an axis is named ``"pressure"``,
+    ``evaluate`` is called with ``pressure=...`` and the map is later loaded
+    with ``inputs={"pressure": some_state}``. Output names come from either the
+    ``outputs`` argument or the keys returned by ``evaluate``.
+
+    Parameters
+    ----------
+    filename:
+        Output HDF5 filename. The ``.h5`` extension is added when omitted.
+    axes:
+        Independent variables to sweep. The output dataset shape follows this
+        order. For example, axes with lengths ``(30, 50)`` create outputs with
+        shape ``(30, 50)``.
+    evaluate:
+        Callable invoked as ``evaluate(**inputs)`` at every grid point.
+        ``inputs`` contains one value for each axis plus every entry in
+        ``constants``. It must return a flat ``dict[str, scalar_number]``.
+        Tuples, lists, arrays, nested dictionaries, strings, and booleans are
+        rejected as map outputs. Return separate scalar keys instead, such as
+        ``{"Y_H2O": ..., "Y_CO2": ...}``.
+    group:
+        HDF5 group name for this map, such as ``"products_tp"``.
+    outputs:
+        Optional list of output names to require and store. If omitted, output
+        names are discovered from the first successful ``evaluate`` call. If
+        provided, each listed name must be returned by ``evaluate``; extra
+        returned keys are ignored.
+    constants:
+        Fixed keyword arguments passed to every ``evaluate`` call but not swept
+        as axes. Constants are stored in the map group's HDF5 attributes, not as
+        ``/axes`` datasets. Use constants for values like a fixed mixture ratio.
+    metadata:
+        Extra user metadata stored as JSON in the map group's attributes.
+    resume:
+        If ``True`` and the map group already exists, reuse successful grid
+        points from ``status/success`` and only evaluate missing or failed
+        points. The existing axes and requested outputs must match.
+    overwrite:
+        If ``True`` and the map group already exists, delete that group and
+        regenerate it from scratch. This takes priority over ``resume``. Use
+        this when you changed axes, outputs, constants, metadata, or the
+        equations inside ``evaluate`` and want a clean file.
+    fill_value:
+        Value written to output datasets for failed points when
+        ``raise_errors=False``. The default is ``np.nan``.
+    compression, compression_opts:
+        HDF5 compression options for output and status datasets.
+    flush_every:
+        Flush the file to disk after this many evaluated points. Use a smaller
+        value for safer long runs or a larger value for slightly faster writes.
+    raise_errors:
+        If ``True``, immediately raise the original exception from
+        ``evaluate`` and stop generation. If ``False``, failed points are
+        recorded in ``status/message``, ``status/success`` is set to ``False``,
+        and outputs are filled with ``fill_value``.
+
+    Returns
+    -------
+    str
+        Normalized HDF5 filename that was written.
+    """
+    filename = hdf5_filename(filename)
+    axes = list(axes)
+    constants = {} if constants is None else dict(constants)
+    metadata = {} if metadata is None else dict(metadata)
+
+    if not axes:
+        raise ValueError("At least one axis is required.")
+
+    axis_names = [axis.name for axis in axes]
+
+    if len(axis_names) != len(set(axis_names)):
+        raise ValueError("Axis names must be unique.")
+
+    if set(axis_names) & set(constants):
+        raise ValueError("Axis names cannot also be used as constant input names.")
+
+    for axis in axes:
+        _validate_axis(axis)
+
+    if outputs is not None:
+        outputs = list(outputs)
+
+        if not outputs:
+            raise ValueError("outputs must contain at least one output name when provided.")
+
+        if len(outputs) != len(set(outputs)):
+            raise ValueError("Output names must be unique.")
+
+        for output_name in outputs:
+            _validate_hdf5_name(output_name, "Output")
+
+            if output_name in _RESERVED_GROUP_NAMES:
+                raise ValueError(f"Output name '{output_name}' is reserved.")
+
+    group = group.strip("/") or "map"
+    group_path = safe_group_name(group)
+
+    with h5py.File(filename, "a") as file:
+        file.attrs["fullflow_format"] = "fullflow-simple-hdf5"
+        file.attrs["fullflow_schema_version"] = 3
+        if "created_utc" not in file.attrs:
+            file.attrs["created_utc"] = _datetime.datetime.now(_datetime.UTC).isoformat()
+        file.attrs["updated_utc"] = _datetime.datetime.now(_datetime.UTC).isoformat()
+        if group_path in file:
+            if overwrite:
+                del file[group_path]
+            elif not resume:
+                raise FileExistsError(
+                    f"HDF5 group '{group_path}' already exists in '{filename}'. "
+                    "Use resume=True or overwrite=True."
+                )
+
+        if group_path in file:
+            map_group = file[group_path]
+            output_names = _check_existing_group(map_group, axes, outputs, constants, metadata)
+
+        else:
+            if outputs is None:
+                output_names, first_index, first_values, failures = _discover_outputs(
+                    axes,
+                    constants,
+                    evaluate,
+                    outputs,
+                    raise_errors,
+                )
+            else:
+                output_names = list(outputs)
+                first_index = None
+                first_values = None
+                failures = []
+
+            map_group = _create_group(
+                file,
+                group,
+                axes,
+                output_names,
+                constants,
+                metadata,
+                fill_value,
+                compression,
+                compression_opts,
+            )
+
+            for index, failure_message in failures:
+                _write_failure(map_group, index, output_names, fill_value, failure_message)
+
+            if first_index is not None and first_values is not None:
+                for output_name, value in first_values.items():
+                    map_group["outputs"][output_name][first_index] = value
+
+                map_group["status/success"][first_index] = True
+                map_group["status/message"][first_index] = ""
+
+        shape = _shape_from_axes(axes)
+        status_group = map_group["status"]
+        outputs_group = map_group["outputs"]
+        counter = 0
+
+        for index in np.ndindex(shape):
+            if resume and bool(status_group["success"][index]):
+                continue
+
+            inputs = _inputs_from_index(axes, index, constants)
+
+            try:
+                values = _validate_scalar_outputs(evaluate(**inputs), output_names)
+
+                for output_name, value in values.items():
+                    outputs_group[output_name][index] = value
+
+                status_group["success"][index] = True
+                status_group["message"][index] = ""
+
+            except MapOutputError:
+                raise
+            except Exception as exc:
+                if raise_errors:
+                    raise
+
+                _write_failure(
+                    map_group,
+                    index,
+                    output_names,
+                    fill_value,
+                    f"{type(exc).__name__}: {exc}",
+                )
+
+            counter += 1
+
+            if flush_every and counter % flush_every == 0:
+                file.flush()
+
+        file.flush()
+
+    return filename