spacecore 0.1.1__py3-none-any.whl

spacecore/__init__.py

@@ -0,0 +1,46 @@
+__version__ = "0.1.1"
+
+
+from .backend import Context, BackendOps, JaxOps, NumpyOps, jax_pytree_class
+from .linop import DenseLinOp, SparseLinOp, BlockDiagonalLinOp, SumToSingleLinOp, StackedLinOp
+from .space import VectorSpace, HermitianSpace, Space, ProductSpace
+from .types import DenseArray, SparseArray, ArrayLike
+
+from ._contextual.manager import (
+    set_context, get_context,
+    register_ops,
+    set_resolution_policy, set_dtype_resolution_policy,
+    get_resolution_policy, get_dtype_resolution_policy
+)
+
+__all__ = [
+    "Context",
+
+    "BackendOps",
+    "JaxOps",
+    "jax_pytree_class",
+    "NumpyOps",
+
+    "DenseLinOp",
+    "SparseLinOp",
+    "BlockDiagonalLinOp",
+    "SumToSingleLinOp",
+    "StackedLinOp",
+
+    "VectorSpace",
+    "HermitianSpace",
+    "ProductSpace",
+    "Space",
+
+    "DenseArray",
+    "SparseArray",
+    "ArrayLike",
+
+    "set_context",
+    "get_context",
+    "register_ops",
+    "set_resolution_policy",
+    "set_dtype_resolution_policy",
+    "get_resolution_policy",
+    "get_dtype_resolution_policy",
+]

spacecore/_contextual/__init__.py

@@ -0,0 +1,10 @@
+from .bound import ContextBound as ContextBound
+from .manager import (
+    ctx_manager as ctx_manager,
+    set_context as set_context,
+    register_ops as register_ops,
+    set_resolution_policy as set_resolution_policy,
+    set_dtype_resolution_policy as set_dtype_resolution_policy,
+    get_resolution_policy as get_resolution_policy,
+    get_dtype_resolution_policy as get_dtype_resolution_policy,
+)

spacecore/_contextual/bound.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from abc import ABC
+from typing import Self
+
+from ..backend import Context, BackendOps, BackendFamily
+from ..types import DType
+from .manager import ctx_manager
+
+class ContextBound(ABC):
+    def __init__(self, ctx: Context | str | None = None):
+        ctx = ctx_manager.normalize_context(ctx)
+        self._ctx = ctx
+
+    @property
+    def ops(self) -> BackendOps:
+        return self.ctx.ops
+
+    @property
+    def dtype(self) -> DType:
+        return self.ctx.dtype
+
+    @property
+    def ctx(self) -> Context:
+        return self._ctx
+
+    def _convert(self, new_ctx: Context) -> Self:
+        raise NotImplementedError()
+
+    def convert(self, new_ctx: Context | BackendFamily | str | None = None) -> Self:
+        _, new_ctx = ctx_manager.enforce_convert_policy(self, new_ctx)
+        return self._convert(new_ctx)

spacecore/_contextual/contextual.py

@@ -0,0 +1,354 @@
+from __future__ import annotations
+
+from typing import Dict, Any, Iterable, Tuple
+from enum import StrEnum, auto
+from warnings import warn
+
+from ..types import DType
+from ..backend import Context, NumpyOps, JaxOps, BackendFamily, BackendOps
+
+
+class ContextPolicy(StrEnum):
+    warning = auto()
+    error = auto()
+    silent = auto()
+
+class DtypePreservePolicy(StrEnum):
+    keep_native = auto()
+    convert = auto()
+
+
+class ContextError(RuntimeError):
+    pass
+
+
+class ContextInferenceError(ContextError):
+    pass
+
+
+class ContextConflictError(ContextError):
+    pass
+
+
+class UnknownBackendError(ContextError):
+    pass
+
+class ContextConversionError(ContextError):
+    pass
+
+
+class Contextual:
+    """
+    Backend resolver.
+    """
+    _default_ctx: Context
+    _available_ops: Dict[str, type[BackendOps]]
+    _resolution_policy: ContextPolicy
+
+    _default_policy: ContextPolicy = ContextPolicy.warning
+    _default_dtype_resolution_policy: DtypePreservePolicy = DtypePreservePolicy.keep_native
+    _default_dtype: DType | None = None
+    _default_enable_checks: bool = False
+
+    def __init__(self,
+                 resolution_policy: str | ContextPolicy | None = None,
+                 dtype_resolution_policy: str | DtypePreservePolicy | None = None
+                 ) -> None:
+        ops = NumpyOps()
+        self.default_ctx = Context(
+            ops=ops,
+            dtype=ops.sanitize_dtype(self._default_dtype),
+            enable_checks=self._default_enable_checks
+        )
+
+        self._available_ops = {
+            self._backend_key(NumpyOps): NumpyOps,
+            self._backend_key(JaxOps): JaxOps,
+        }
+
+        self.resolution_policy = resolution_policy
+        self.dtype_resolution_policy = dtype_resolution_policy
+
+    def normalize_context(self,
+                          ctx: Context | BackendFamily | str | None = None,
+                          dtype: Any = None,
+                          enable_checks: bool | None = None
+                          ) -> Context:
+        if ctx is None:
+            if dtype is not None or enable_checks is not None:
+                warn(
+                    'Provided context is None, dtype and enable_checks parameters are ignored.',
+                    UserWarning,
+                )
+            return self.default_ctx
+        if isinstance(ctx, Context):
+            if dtype is not None or enable_checks is not None:
+                warn(
+                    'Provided concrete context, dtype and enable_checks parameters are ignored.',
+                    UserWarning,
+                )
+            return Context(
+                ops=ctx.ops,
+                dtype=ctx.ops.sanitize_dtype(ctx.dtype),
+                enable_checks=ctx.enable_checks
+            )
+        if isinstance(ctx, (str, BackendFamily)):
+            ctx = self._backend_key(ctx)
+            ops = self.get_ops(ctx)
+            return self.ctx_from_ops(ops, dtype=dtype, enable_checks=enable_checks)
+        else:
+            raise TypeError(f'Expected Context, BackendFamily, str, or None, got {type(ctx)}.')
+
+    def ctx_like(self, base: Context | None, ctx: Context) -> Context:
+        if isinstance(base, Context):
+            return Context(
+                ops=ctx.ops,
+                dtype=ctx.ops.sanitize_dtype(base.dtype),
+                enable_checks=ctx.enable_checks,
+            )
+        return self.default_ctx
+
+    def normalize_context_like(self, base: Context | None, ctx: Context | BackendFamily | str | None = None) -> Context:
+        if self.dtype_resolution_policy is DtypePreservePolicy.keep_native and isinstance(base, Context):
+            if ctx is None:
+                return self.ctx_like(base, self.default_ctx)
+            if isinstance(ctx, Context):
+                return self.ctx_like(base, ctx)
+            if isinstance(ctx, (str, BackendFamily)):
+                ctx = self._backend_key(ctx)
+                ops = self.get_ops(ctx)
+                return self.ctx_from_ops(
+                    ops=ops,
+                    dtype=base.dtype,
+                    enable_checks=self._default_enable_checks,
+                )
+            else:
+                raise TypeError(f'Expected Context, BackendFamily, str, or None, got {type(ctx)}.')
+        else:
+            return self.normalize_context(ctx)
+
+    def ctx_from_ops(self, ops: BackendOps, dtype: DType | None = None, enable_checks: bool | None = None) -> Context:
+        dtype = ops.sanitize_dtype(dtype)
+        if enable_checks is None:
+            enable_checks = self._default_enable_checks
+        return Context(ops=ops,
+                       dtype=dtype,
+                       enable_checks=enable_checks)
+
+    @property
+    def default_ctx(self) -> Context:
+        return self._default_ctx
+
+    @default_ctx.setter
+    def default_ctx(self, ctx: Context | BackendFamily | str | None = None) -> None:
+        ctx = self.normalize_context(ctx)
+        self._default_ctx = ctx
+
+    @property
+    def resolution_policy(self) -> ContextPolicy:
+        return self._resolution_policy
+
+    @resolution_policy.setter
+    def resolution_policy(self, policy: str | None = ContextPolicy.warning.value) -> None:
+        if policy is None:
+            self._resolution_policy = self._default_policy
+            return
+
+        try:
+            self._resolution_policy = (
+                policy
+                if isinstance(policy, ContextPolicy)
+                else ContextPolicy(policy)
+            )
+        except ValueError as e:
+            allowed = ", ".join(p.value for p in ContextPolicy)
+            raise ValueError(
+                f"Unknown resolution_policy={policy!r}. "
+                f"Expected one of: {allowed}"
+            ) from e
+
+    @property
+    def dtype_resolution_policy(self) -> DtypePreservePolicy:
+        return self._dtype_resolution_policy
+
+    @dtype_resolution_policy.setter
+    def dtype_resolution_policy(self, policy: str | None = DtypePreservePolicy.keep_native.value) -> None:
+        if policy is None:
+            self._dtype_resolution_policy = self._default_dtype_resolution_policy
+            return
+
+        try:
+            self._dtype_resolution_policy = (
+                policy
+                if isinstance(policy, DtypePreservePolicy)
+                else DtypePreservePolicy(policy)
+            )
+        except ValueError as e:
+            allowed = ", ".join(p.value for p in DtypePreservePolicy)
+            raise ValueError(
+                f"Unknown dtype_resolution_policy={policy!r}. "
+                f"Expected one of: {allowed}"
+            ) from e
+
+    def get_ops(self, name: str | BackendFamily | BackendOps | type[BackendOps] | Context) -> BackendOps:
+        name = self._backend_key(name)
+        if name not in self.available_ops:
+            allowed = ", ".join(k for k in self.available_ops.keys())
+            raise UnknownBackendError(
+                f"Unknown backend: {name!r}. "
+                f"Expected one of: {allowed}"
+            )
+        return self.available_ops[name]()
+
+    @property
+    def available_ops(self) -> Dict[str, type[BackendOps]]:
+        return self._available_ops
+
+    def register_ops(self, ops: type[BackendOps]) -> type[BackendOps]:
+        if not isinstance(ops, type) or not issubclass(ops, BackendOps):
+            raise TypeError(f"Expected type[BackendOps], got {type(ops)!r}")
+        else:
+            family = self._backend_key(ops)
+            if family in self.available_ops.keys():
+                raise ContextConflictError(f'BackendOps {family} is already registered.')
+            self._available_ops[family] = ops
+            return ops
+
+    def infer_context(self, x: Any, enable_checks: bool | None = None) -> Context | None:
+        """
+        Infer a context from a single value.
+
+        Intended precedence:
+          1. objects carrying `.ctx`
+          2. backend-native arrays recognized by registered backends
+        """
+        if isinstance(x, Context):
+            return x
+
+        ctx = getattr(x, "ctx", None)
+        if isinstance(ctx, Context):
+            return ctx
+
+        matched: list[BackendOps] = []
+        for name, ops in self.available_ops.items():
+            try:
+                ops = ops()
+                if ops.is_array(x):
+                    matched.append(ops)
+            except Exception:
+                # Keep inference conservative.
+                continue
+
+        if not matched:
+            return None
+        if len(matched) > 1:
+            raise ContextInferenceError(
+                f"Ambiguous backend inference for object of type {type(x)!r}: {matched!r}."
+            )
+
+        ops = matched[0]
+        try:
+            dtype = ops.get_dtype(x)
+        except Exception:
+            dtype = getattr(x, "dtype", self.default_ctx.dtype)
+
+        return self.ctx_from_ops(ops, dtype, enable_checks)
+
+    def infer_contexts(self, values: Iterable[Any]) -> Tuple[Context, ...]:
+        out: list[Context] = []
+        for x in values:
+            ctx = self.infer_context(x)
+            if ctx is not None:
+                out.append(ctx)
+        return tuple(out)
+
+    def are_compatible_contexts(self, *ctxs: Context) -> bool:
+        if len(ctxs) < 2:
+            return True
+        first = ctxs[0]
+        return all(ctx.ops.family == first.ops.family for ctx in ctxs[1:])
+
+    def are_compatible_values(self, *values: Any) -> bool:
+        return self.are_compatible_contexts(*self.infer_contexts(values))
+
+    def are_compatible_ops(self, *ops: BackendOps) -> bool:
+        if not ops:
+            return True
+        first = ops[0]
+        return all(op.family == first.family for op in ops)
+
+    def enforce_convert_policy(self, x: Any, to: Context | BackendFamily | str | None = None) -> Tuple[Any, Context]:
+        native_ctx = self.infer_context(x)
+        ctx = self.normalize_context_like(native_ctx, to)
+        if self.resolution_policy is not ContextPolicy.silent:
+            if native_ctx is not None and not self.are_compatible_contexts(native_ctx, ctx):
+                if self.resolution_policy is ContextPolicy.warning:
+                    warn(
+                        f"Converting from {native_ctx!r} to {ctx!r}.",
+                        UserWarning,
+                    )
+                else:
+                    raise ContextConversionError(
+                        f"Conversion from {native_ctx!r} to {ctx!r} is forbidden by policy {self.resolution_policy.value!r}."
+                    )
+        return x, ctx
+
+    def _backend_key(self, x: str | BackendFamily | BackendOps | type[BackendOps] | Context) -> str:
+        if isinstance(x, Context):
+            return self._backend_key(x.ops)
+        if isinstance(x, BackendOps):
+            return self._backend_key(x.family)
+        if isinstance(x, type) and issubclass(x, BackendOps):
+            return self._backend_key(x._family)
+        if isinstance(x, BackendFamily):
+            return x.value.lower()
+        if isinstance(x, str):
+            return x.lower()
+        raise TypeError(f"Unsupported backend key source: {type(x)!r}")
+
+    def resolve_context_priority(
+            self,
+            priority_ctx: Context | BackendFamily | str | None = None,
+            *other_ctx: object,
+    ) -> Context:
+        """
+        Resolve context with priority:
+
+        1. If priority_ctx is not None, return its normalized form.
+        2. Otherwise, infer contexts from other_ctx and return the best inferred one.
+        3. If inference fails, return default_ctx.
+
+        Policy:
+        - explicit context always wins
+        - inferred contexts must be family-compatible
+        - if inferred dtypes differ, take the most general of them
+        """
+        if priority_ctx is not None:
+            return self.normalize_context(priority_ctx)
+
+        inferred = self.infer_contexts(other_ctx)
+        if not inferred:
+            return self.default_ctx
+
+        if not self.are_compatible_contexts(*inferred):
+            fams = tuple(ctx.ops.family for ctx in inferred)
+            raise ValueError(f"Incompatible inferred contexts: {fams!r}")
+
+        first = inferred[0]
+        ops = type(first.ops)()
+        dtype = self._join_dtypes(ops, *(ctx.dtype for ctx in inferred))
+
+        return self.ctx_from_ops(
+            ops=ops,
+            dtype=dtype,
+            enable_checks=all(ctx.enable_checks for ctx in inferred),
+        )
+
+    def _join_dtypes(self, ops: BackendOps, *dtypes: DType | None) -> DType | None:
+        clean = [ops.sanitize_dtype(dt) for dt in dtypes if dt is not None]
+        if not clean:
+            return ops.sanitize_dtype(None)
+
+        np_ops = NumpyOps()
+        joined = np_ops.np.result_type(*clean)
+        return ops.sanitize_dtype(joined)

spacecore/_contextual/manager.py

@@ -0,0 +1,43 @@
+from typing import Any
+
+from ..backend import Context, BackendOps
+from .contextual import Contextual, ContextPolicy, DtypePreservePolicy
+from ..backend import BackendFamily
+
+
+ctx_manager = Contextual()
+
+
+def set_context(
+        ctx: Context | BackendFamily | str | None = None,
+        dtype: Any = None,
+        enable_checks: bool | None = None
+) -> None:
+    ctx = ctx_manager.normalize_context(ctx, dtype=dtype, enable_checks=enable_checks)
+    ctx_manager.default_ctx = ctx
+
+
+def get_context():
+    return ctx_manager.default_ctx
+
+
+def register_ops(ops: type[BackendOps]) -> type[BackendOps]:
+    return ctx_manager.register_ops(ops)
+
+
+def set_resolution_policy(policy: ContextPolicy | str | None = None) -> None:
+    ctx_manager.resolution_policy = policy
+
+
+def get_resolution_policy() -> str:
+    return ctx_manager.resolution_policy.value
+
+
+def set_dtype_resolution_policy(
+    policy: DtypePreservePolicy | str | None = None,
+) -> None:
+    ctx_manager.dtype_resolution_policy = policy
+
+
+def get_dtype_resolution_policy() -> str:
+    return ctx_manager.dtype_resolution_policy.value

spacecore/backend/__init__.py

@@ -0,0 +1,15 @@
+from ._context import Context
+from ._ops import BackendOps
+from ._family import BackendFamily
+from .jax import JaxOps, jax_pytree_class
+from .numpy import NumpyOps
+
+
+__all__ = [
+    "Context",
+    "BackendFamily",
+    "BackendOps",
+    "JaxOps",
+    "jax_pytree_class",
+    "NumpyOps",
+]

spacecore/backend/_context.py

@@ -0,0 +1,50 @@
+from dataclasses import dataclass
+from typing import Any
+
+from ._ops import BackendOps
+from ..types import DenseArray, SparseArray, DType, ArrayLike
+
+
+@dataclass(frozen=True, slots=True)
+class Context:
+    ops: BackendOps
+    dtype: DType | None = None
+    enable_checks: bool = True
+
+    def __post_init__(self):
+        if not isinstance(self.ops, BackendOps):
+            raise TypeError("ops must be a BackendOps")
+
+        sanitized = self.ops.sanitize_dtype(self.dtype)
+        object.__setattr__(self, "dtype", sanitized)
+
+    def assert_dense(self, x: Any) -> DenseArray:
+        if not self.ops.is_dense(x):
+            raise TypeError(f"Expected dense array for {self.ops.family}, got {type(x).__name__}")
+        return x
+
+    def assert_sparse(self, x: Any) -> SparseArray:
+        if not self.ops.allow_sparse:
+            raise TypeError("Sparse objects are disallowed by this backend.")
+        if not self.ops.is_sparse(x):
+            raise TypeError(f"Expected sparse array for {self.ops.family}, got {type(x).__name__}")
+        return x
+
+    def asarray(self, x: Any) -> DenseArray:
+        return self.ops.asarray(x, dtype=self.dtype)
+
+    def assparse(self, x: Any) -> SparseArray:
+        return self.ops.assparse(x, dtype=self.dtype)
+
+    def convert(self, x: Any) -> ArrayLike:
+        if self.ops.is_dense(x):
+            return self.asarray(x)
+        elif self.ops.is_sparse(x):
+            return self.assparse(x)
+        else:
+            raise NotImplementedError
+
+    def __eq__(self, other: Any) -> bool:
+        if isinstance(other, Context):
+            return self.ops == other.ops and self.enable_checks == other.enable_checks
+        return False

spacecore/backend/_family.py

@@ -0,0 +1,6 @@
+from enum import StrEnum, auto
+
+
+class BackendFamily(StrEnum):
+    numpy = auto()
+    jax = auto()