optixde 0.2.2__py3-none-any.whl

optixde/__init__.py

@@ -0,0 +1,10 @@
+# SPDX-FileCopyrightText: 2025-2026 Yang Yang <yangyhhu@foxmail.com>
+# SPDX-License-Identifier: MIT
+
+
+from .version import __version__
+from . import geometry, bc, solvers, post, operators, fft_backend
+
+__all__ = [
+    '__version__', 'geometry', 'bc', 'solvers', 'post', 'operators', 'fft_backend'
+]

optixde/bc/__init__.py

@@ -0,0 +1,25 @@
+# SPDX-FileCopyrightText: 2025-2026 Yang Yang <yangyhhu@foxmail.com>
+# SPDX-License-Identifier: MIT
+
+
+"""Boundary-condition helpers.
+
+This subpackage centralizes utilities for enforcing boundary conditions in a
+backend-aware way (NumPy / CuPy / PyTorch), and also exposes a segmented
+polygon-boundary API for mixed boundary conditions on the same geometric edge.
+"""
+
+from .robin import normalize_robin_spec, apply_robin_fd_projection
+from .segmented import EdgeSlice, DirichletBC, NeumannBC, RobinBC, split_bcs_by_type
+from .rasterize import rasterize_segmented_bcs
+
+__all__ = [
+    "normalize_robin_spec",
+    "apply_robin_fd_projection",
+    "EdgeSlice",
+    "DirichletBC",
+    "NeumannBC",
+    "RobinBC",
+    "split_bcs_by_type",
+    "rasterize_segmented_bcs",
+]

optixde/bc/rasterize.py

@@ -0,0 +1,100 @@
+# SPDX-FileCopyrightText: 2025-2026 Yang Yang <yangyhhu@foxmail.com>
+# SPDX-License-Identifier: MIT
+
+
+from __future__ import annotations
+
+import numpy as np
+
+from .segmented import DirichletBC, NeumannBC, RobinBC
+
+
+def rasterize_segmented_bcs(geometry, X, bcs, band_width):
+    """Rasterize segmented edge BCs onto query points.
+
+    Parameters
+    ----------
+    geometry : Polygon-like
+        Must implement `project_to_edges(X)` and `edge_outward_normal(edge_id)`.
+    X : ndarray, shape (N, 2)
+        Query points.
+    bcs : list
+        Boundary condition objects from optixde.bc.segmented.
+    band_width : float
+        Boundary activation half-width.
+
+    Returns
+    -------
+    dict of ndarrays
+        type_id: 0 none, 1 dirichlet, 2 neumann, 3 robin
+        distance, nx, ny, value, alpha, beta, weight
+    """
+    X = np.asarray(X, dtype=float)
+    N = X.shape[0]
+    out = {
+        'type_id': np.zeros(N, dtype=int),
+        'distance': np.full(N, np.inf, dtype=float),
+        'nx': np.zeros(N, dtype=float),
+        'ny': np.zeros(N, dtype=float),
+        'value': np.zeros(N, dtype=float),
+        'alpha': np.zeros(N, dtype=float),
+        'beta': np.zeros(N, dtype=float),
+        'weight': np.zeros(N, dtype=float),
+        'edge_id': np.full(N, -1, dtype=int),
+        's': np.zeros(N, dtype=float),
+    }
+    if len(bcs) == 0:
+        return out
+
+    proj = geometry.project_to_edges(X)
+    edges = proj['edge_id']
+    svals = proj['s']
+    dists = proj['distance']
+
+    for bc in sorted(bcs, key=lambda item: getattr(item, 'priority', 0), reverse=True):
+        seg = bc.segment
+        hit = (
+            (edges == int(seg.edge_id))
+            & (svals >= float(seg.s0) - 1e-12)
+            & (svals <= float(seg.s1) + 1e-12)
+            & (dists <= float(band_width))
+        )
+        if not np.any(hit):
+            continue
+
+        replace = hit & (
+            (dists < out['distance'] - 1e-15)
+            | ((np.abs(dists - out['distance']) <= 1e-15) & (getattr(bc, 'priority', 0) >= 0))
+        )
+        if not np.any(replace):
+            continue
+
+        normal = geometry.edge_outward_normal(seg.edge_id)
+        x = X[replace, 0]
+        y = X[replace, 1]
+        out['distance'][replace] = dists[replace]
+        out['nx'][replace] = normal[0]
+        out['ny'][replace] = normal[1]
+        out['edge_id'][replace] = int(seg.edge_id)
+        out['s'][replace] = svals[replace]
+        out['weight'][replace] = geometry.smooth_dirac(dists[replace], band_width)
+
+        if isinstance(bc, DirichletBC):
+            out['type_id'][replace] = 1
+            out['value'][replace] = bc.value_callable()(x, y)
+            out['alpha'][replace] = 1.0
+            out['beta'][replace] = 0.0
+        elif isinstance(bc, NeumannBC):
+            out['type_id'][replace] = 2
+            out['value'][replace] = bc.flux_callable()(x, y)
+            out['alpha'][replace] = 0.0
+            out['beta'][replace] = 1.0
+        elif isinstance(bc, RobinBC):
+            out['type_id'][replace] = 3
+            out['value'][replace] = bc.value_callable()(x, y)
+            out['alpha'][replace] = bc.alpha_callable()(x, y)
+            out['beta'][replace] = bc.beta_callable()(x, y)
+        else:
+            raise TypeError(f'Unsupported segmented BC type: {type(bc)!r}')
+
+    return out

optixde/bc/robin.py

@@ -0,0 +1,149 @@
+# SPDX-FileCopyrightText: 2025-2026 Yang Yang <yangyhhu@foxmail.com>
+# SPDX-License-Identifier: MIT
+
+"""Robin (mixed) boundary condition helpers.
+
+Convention on each side:
+
+    alpha * u + beta * (∂u/∂n) = g
+
+For rectangular grids that include boundary nodes, we enforce Robin conditions
+via a one-sided finite-difference projection.
+
+This implementation is backend-aware: NumPy / CuPy / PyTorch.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Tuple, Union
+
+
+def _is_cupy_array(x: Any) -> bool:
+    try:
+        import cupy as cp
+        return isinstance(x, cp.ndarray)
+    except Exception:
+        return False
+
+
+def _is_torch_tensor(x: Any) -> bool:
+    try:
+        import torch
+        return isinstance(x, torch.Tensor)
+    except Exception:
+        return False
+
+
+def array_namespace(x: Any):
+    if _is_torch_tensor(x):
+        import torch
+        return torch
+    if _is_cupy_array(x):
+        import cupy as cp
+        return cp
+    import numpy as np
+    return np
+
+
+RobinSide = Tuple[float, float, Any]  # (alpha, beta, g)
+RobinSpec = Dict[str, RobinSide]
+RobinInput = Union[None, RobinSide, Dict[str, RobinSide]]
+
+
+def _as_side_value(xp, g: Any, n: int, ref):
+    if g is None:
+        if xp.__name__ == "torch":
+            return xp.zeros((n,), dtype=ref.dtype, device=ref.device)
+        return xp.zeros((n,), dtype=float)
+
+    if xp.__name__ == "torch":
+        import torch
+
+        if isinstance(g, torch.Tensor):
+            if g.ndim == 0:
+                return g.expand(n)
+            if g.shape[0] != n:
+                raise ValueError(f"Robin g has length {g.shape[0]}, expected {n}.")
+            return g
+        if isinstance(g, (int, float)):
+            return torch.full((n,), float(g), dtype=ref.dtype, device=ref.device)
+        gg = torch.as_tensor(g, dtype=ref.dtype, device=ref.device)
+        if gg.ndim == 0:
+            return gg.expand(n)
+        if gg.shape[0] != n:
+            raise ValueError(f"Robin g has length {gg.shape[0]}, expected {n}.")
+        return gg
+
+    if xp.isscalar(g):
+        return xp.full((n,), float(g), dtype=float)
+    arr = xp.asarray(g)
+    if arr.ndim == 0:
+        return xp.full((n,), float(arr), dtype=float)
+    if arr.shape[0] != n:
+        raise ValueError(f"Robin g has length {arr.shape[0]}, expected {n}.")
+    return arr.astype(float)
+
+
+def normalize_robin_spec(robin: RobinInput) -> RobinSpec:
+    default: RobinSide = (1.0, 0.0, 0.0)
+    if robin is None:
+        return {"left": default, "right": default, "bottom": default, "top": default}
+    if isinstance(robin, tuple) and len(robin) == 3:
+        a, b, g = robin
+        return {"left": (a, b, g), "right": (a, b, g), "bottom": (a, b, g), "top": (a, b, g)}
+    if isinstance(robin, dict):
+        base = robin.get("default", default)
+        out: RobinSpec = {}
+        for side in ("left", "right", "bottom", "top"):
+            out[side] = robin.get(side, base)
+            if not (isinstance(out[side], tuple) and len(out[side]) == 3):
+                raise ValueError(
+                    "Each Robin side must be a tuple (alpha, beta, g). "
+                    f"Got {side}={out[side]!r}."
+                )
+        return out
+    raise ValueError(
+        "Robin must be None, a tuple (alpha,beta,g), or a dict with sides. "
+        f"Got {type(robin)}."
+    )
+
+
+def apply_robin_fd_projection(u, Lx: float, Ly: float, robin: RobinInput) -> Any:
+    xp = array_namespace(u)
+    Ny, Nx = int(u.shape[0]), int(u.shape[1])
+    if Nx < 2 or Ny < 2:
+        raise ValueError("Robin projection requires at least a 2×2 grid including boundary nodes.")
+
+    hx = float(Lx) / float(Nx - 1)
+    hy = float(Ly) / float(Ny - 1)
+    spec = normalize_robin_spec(robin)
+
+    a, b, g = spec["left"]
+    g1 = _as_side_value(xp, g, Ny, u)
+    denom = float(a) + float(b) / hx
+    if abs(denom) < 1e-30:
+        raise ValueError("Invalid Robin coefficients on left: alpha + beta/h ~ 0.")
+    u[:, 0] = (g1 + (float(b) / hx) * u[:, 1]) / denom
+
+    a, b, g = spec["right"]
+    g1 = _as_side_value(xp, g, Ny, u)
+    denom = float(a) + float(b) / hx
+    if abs(denom) < 1e-30:
+        raise ValueError("Invalid Robin coefficients on right: alpha + beta/h ~ 0.")
+    u[:, -1] = (g1 + (float(b) / hx) * u[:, -2]) / denom
+
+    a, b, g = spec["bottom"]
+    g1 = _as_side_value(xp, g, Nx, u)
+    denom = float(a) + float(b) / hy
+    if abs(denom) < 1e-30:
+        raise ValueError("Invalid Robin coefficients on bottom: alpha + beta/h ~ 0.")
+    u[0, :] = (g1 + (float(b) / hy) * u[1, :]) / denom
+
+    a, b, g = spec["top"]
+    g1 = _as_side_value(xp, g, Nx, u)
+    denom = float(a) + float(b) / hy
+    if abs(denom) < 1e-30:
+        raise ValueError("Invalid Robin coefficients on top: alpha + beta/h ~ 0.")
+    u[-1, :] = (g1 + (float(b) / hy) * u[-2, :]) / denom
+
+    return u

optixde/bc/segmented.py

@@ -0,0 +1,228 @@
+# SPDX-FileCopyrightText: 2025-2026 Yang Yang <yangyhhu@foxmail.com>
+# SPDX-License-Identifier: MIT
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable, Iterable, Optional, Sequence, Tuple, Union
+
+import numpy as np
+
+Number = Union[float, int]
+ScalarOrCallable = Union[Number, Callable[[float, float], Union[Number, np.ndarray]]]
+
+
+
+def _as_callable_2d(value: ScalarOrCallable) -> Callable[[np.ndarray, np.ndarray], np.ndarray]:
+    """Return a vectorized callable f(x, y).
+
+    This preserves the previous user-facing convenience that boundary values may be
+    given either as constants or as callables.
+    """
+    if callable(value):
+        def _wrapped(x, y):
+            out = value(x, y)
+            if np.isscalar(out):
+                return np.full_like(np.asarray(x, dtype=float), float(out), dtype=float)
+            return np.asarray(out, dtype=float)
+        return _wrapped
+
+    def _const(x, y):
+        return np.full_like(np.asarray(x, dtype=float), float(value), dtype=float)
+
+    return _const
+
+
+@dataclass(frozen=True)
+class EdgeSegment:
+    """A sub-segment on a polygon edge.
+
+    Parameters
+    ----------
+    edge_id : int
+        Polygon edge index.
+    s0, s1 : float
+        Local edge coordinates in [0, 1].
+    name : str | None
+        Optional user label retained for backward compatibility.
+    """
+
+    edge_id: int
+    s0: float = 0.0
+    s1: float = 1.0
+    name: Optional[str] = None
+
+    def __post_init__(self):
+        s0 = float(self.s0)
+        s1 = float(self.s1)
+        if s0 > s1:
+            s0, s1 = s1, s0
+        object.__setattr__(self, "s0", max(0.0, min(1.0, s0)))
+        object.__setattr__(self, "s1", max(0.0, min(1.0, s1)))
+        object.__setattr__(self, "edge_id", int(self.edge_id))
+
+
+@dataclass(frozen=True)
+class EdgeSlice:
+    """Backward-compatible geometric edge slice description.
+
+    Old user code may still construct boundary slices using explicit segment end
+    points on a named edge. We keep this object unchanged so existing scripts keep
+    working, then convert it to an ``EdgeSegment`` when geometry is available.
+    """
+
+    name: str
+    start: Tuple[float, float]
+    end: Tuple[float, float]
+    kind: str   # "dirichlet", "neumann", or "robin"
+    value: ScalarOrCallable = 0.0
+
+
+@dataclass
+class BoundaryCondition:
+    """Base boundary-condition object for segmented polygon edges."""
+
+    segment: Union[EdgeSegment, EdgeSlice]
+    priority: int = 0
+
+
+@dataclass
+class DirichletBC(BoundaryCondition):
+    value: ScalarOrCallable = 0.0
+
+    def value_callable(self) -> Callable[[np.ndarray, np.ndarray], np.ndarray]:
+        return _as_callable_2d(self.value)
+
+
+@dataclass
+class NeumannBC(BoundaryCondition):
+    flux: ScalarOrCallable = 0.0
+
+    # Backward-compatible alias
+    @property
+    def value(self):
+        return self.flux
+
+    def flux_callable(self) -> Callable[[np.ndarray, np.ndarray], np.ndarray]:
+        return _as_callable_2d(self.flux)
+
+
+@dataclass
+class RobinBC(BoundaryCondition):
+    alpha: ScalarOrCallable = 1.0
+    beta: ScalarOrCallable = 1.0
+    value: ScalarOrCallable = 0.0
+
+    def alpha_callable(self) -> Callable[[np.ndarray, np.ndarray], np.ndarray]:
+        return _as_callable_2d(self.alpha)
+
+    def beta_callable(self) -> Callable[[np.ndarray, np.ndarray], np.ndarray]:
+        return _as_callable_2d(self.beta)
+
+    def value_callable(self) -> Callable[[np.ndarray, np.ndarray], np.ndarray]:
+        return _as_callable_2d(self.value)
+
+
+
+def _point_on_segment_parameter(p0: np.ndarray, p1: np.ndarray, q: np.ndarray) -> float:
+    v = p1 - p0
+    vv = float(np.dot(v, v))
+    if vv <= 0.0:
+        return 0.0
+    return float(np.clip(np.dot(q - p0, v) / vv, 0.0, 1.0))
+
+
+
+def _resolve_edge_slice(geometry, edge_slice: EdgeSlice) -> EdgeSegment:
+    """Map a legacy EdgeSlice onto geometry edge_id and local [s0, s1]."""
+    if geometry is None:
+        raise ValueError("Geometry is required to resolve a legacy EdgeSlice into edge_id/s-range.")
+    if not hasattr(geometry, "project_to_edges"):
+        raise TypeError("Geometry must implement project_to_edges to use legacy EdgeSlice objects.")
+
+    pts = np.asarray([edge_slice.start, edge_slice.end], dtype=float)
+    proj = geometry.project_to_edges(pts)
+    edge_ids = np.asarray(proj["edge_id"], dtype=int)
+    svals = np.asarray(proj["s"], dtype=float)
+    if edge_ids.shape[0] != 2:
+        raise ValueError("Failed to resolve EdgeSlice endpoints on geometry edges.")
+    if edge_ids[0] != edge_ids[1]:
+        raise ValueError(
+            f"EdgeSlice '{edge_slice.name}' spans multiple polygon edges ({edge_ids[0]} and {edge_ids[1]}). "
+            "A single segmented BC must lie on one edge."
+        )
+    return EdgeSegment(edge_id=int(edge_ids[0]), s0=float(min(svals)), s1=float(max(svals)), name=edge_slice.name)
+
+
+
+def normalize_segment(segment: Union[EdgeSegment, EdgeSlice], geometry=None) -> EdgeSegment:
+    if isinstance(segment, EdgeSegment):
+        return segment
+    if isinstance(segment, EdgeSlice):
+        return _resolve_edge_slice(geometry, segment)
+    raise TypeError(f"Unsupported segment type: {type(segment)!r}")
+
+
+
+def normalize_bc(bc, geometry=None):
+    """Return a BC whose ``segment`` is always an ``EdgeSegment``.
+
+    This keeps old scripts working while matching the rasterizer/solver contract.
+    """
+    if isinstance(bc, DirichletBC):
+        return DirichletBC(segment=normalize_segment(bc.segment, geometry), priority=bc.priority, value=bc.value)
+    if isinstance(bc, NeumannBC):
+        return NeumannBC(segment=normalize_segment(bc.segment, geometry), priority=bc.priority, flux=bc.flux)
+    if isinstance(bc, RobinBC):
+        return RobinBC(
+            segment=normalize_segment(bc.segment, geometry),
+            priority=bc.priority,
+            alpha=bc.alpha,
+            beta=bc.beta,
+            value=bc.value,
+        )
+    raise TypeError(f"Unsupported boundary condition type: {type(bc)!r}")
+
+
+
+def normalize_bcs(bcs: Optional[Sequence], geometry=None):
+    if bcs is None:
+        return []
+    return [normalize_bc(bc, geometry) for bc in bcs]
+
+
+
+def split_bcs_by_type(bcs: Iterable):
+    """Split BCs by type while preserving order.
+
+    Returned tuple: (dirichlet, neumann, robin)
+    """
+    dirichlet = []
+    neumann = []
+    robin = []
+    for bc in bcs:
+        if isinstance(bc, DirichletBC):
+            dirichlet.append(bc)
+        elif isinstance(bc, NeumannBC):
+            neumann.append(bc)
+        elif isinstance(bc, RobinBC):
+            robin.append(bc)
+        else:
+            raise TypeError(f"Unsupported boundary condition type: {type(bc)!r}")
+    return dirichlet, neumann, robin
+
+
+__all__ = [
+    "Number",
+    "ScalarOrCallable",
+    "EdgeSegment",
+    "EdgeSlice",
+    "BoundaryCondition",
+    "DirichletBC",
+    "NeumannBC",
+    "RobinBC",
+    "normalize_segment",
+    "normalize_bc",
+    "normalize_bcs",
+    "split_bcs_by_type",
+]

optixde/fft_backend/__init__.py

@@ -0,0 +1,98 @@
+# SPDX-FileCopyrightText: 2025-2026 Yang Yang <yangyhhu@foxmail.com>
+# SPDX-License-Identifier: MIT
+
+"""FFT backend abstraction for OptiXDE.
+
+Backends:
+- NumPyBackend (CPU)
+- TorchBackend (CPU/GPU)
+- CuPyBackend (GPU via cuFFT)  [optional dependency: cupy]
+"""
+
+from .base import FFTBackend, FreqGrids as FreqGrids
+from .numpy_backend import NumpyBackend
+from .propagator_cache import PropagatorCache as PropagatorCache
+
+__all__ = [
+    "FFTBackend",
+    "FreqGrids",
+    "NumpyBackend",
+    "TorchBackend",
+    "CuPyBackend",
+    "PropagatorCache",
+    "get_backend",
+    "clear_backend_cache",
+]
+
+
+_BACKEND_INSTANCES = {}
+
+
+def _backend_key(name, kwargs):
+    """Build a stable key for the small set of backend constructor options."""
+    normalized = []
+    for key, value in sorted(kwargs.items()):
+        try:
+            hash(value)
+            normalized_value = value
+        except TypeError:
+            normalized_value = repr(value)
+        normalized.append((key, normalized_value))
+    return name, tuple(normalized)
+
+
+def clear_backend_cache():
+    """Clear shared backend instances and all arrays cached by them."""
+    for backend in _BACKEND_INSTANCES.values():
+        backend.clear_caches()
+    _BACKEND_INSTANCES.clear()
+
+
+def __getattr__(name):
+    if name == "TorchBackend":
+        try:
+            from .torch_backend import TorchBackend
+        except ModuleNotFoundError as exc:
+            if exc.name == "torch":
+                raise ImportError(
+                    "TorchBackend requires torch. Install PyTorch to use this backend."
+                ) from exc
+            raise
+        return TorchBackend
+
+    if name == "CuPyBackend":
+        try:
+            from .cupy_backend import CuPyBackend
+        except ModuleNotFoundError as exc:
+            if exc.name == "cupy":
+                raise ImportError(
+                    "CuPyBackend requires cupy. Install CuPy to use this backend."
+                ) from exc
+            raise
+        return CuPyBackend
+
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def get_backend(name: str = "numpy", **kwargs) -> FFTBackend:
+    name = (name or "numpy").lower()
+    if name in ("np", "numpy"):
+        canonical_name = "numpy"
+        factory = NumpyBackend
+    if name in ("torch", "pytorch"):
+        from .torch_backend import TorchBackend
+        canonical_name = "torch"
+        factory = TorchBackend
+    elif name in ("cupy", "cp", "cufft"):
+        from .cupy_backend import CuPyBackend
+        canonical_name = "cupy"
+        factory = CuPyBackend
+    elif name not in ("np", "numpy"):
+        raise ValueError(f"Unknown backend: {name!r}. Use 'numpy', 'torch', or 'cupy'.")
+
+    key = _backend_key(canonical_name, kwargs)
+    backend = _BACKEND_INSTANCES.get(key)
+    if backend is None:
+        backend = factory(**kwargs)
+        _BACKEND_INSTANCES[key] = backend
+    return backend