engine/sim/entropy.py

"""noctilith/sim/entropy.py — EntropyField."""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Tuple

import numpy as np


@dataclass(init=False)
class EntropyField:
    """Simple entropy field with backward-compatible constructor."""

    S: np.ndarray = field(init=False)

    def __init__(self, shape_or_array: Any = (1, 1, 1), *, S: Any = None):
        if S is not None:
            arr = np.asarray(S, dtype=np.float64)
        elif isinstance(shape_or_array, np.ndarray):
            arr = np.asarray(shape_or_array, dtype=np.float64)
        else:
            arr = np.zeros(tuple(shape_or_array), dtype=np.float64)
        if arr.ndim != 3:
            raise ValueError("EntropyField expects a 3D array or 3D shape")
        self.S = arr

    @classmethod
    def zeros(cls, shape: Tuple[int, ...]) -> "EntropyField":
        return cls(shape)

    @property
    def shape(self) -> Tuple[int, ...]:
        return self.S.shape

    def add_entropy(self, ix: int, iy: int, iz: int, delta: float) -> None:
        self.S[ix, iy, iz] = float(np.clip(self.S[ix, iy, iz] + delta, 0.0, 1.0))

    def max_entropy(self) -> float:
        return float(np.max(self.S))