scirex 1.0.0__py3-none-any.whl

scirex/__init__.py

@@ -0,0 +1,69 @@
+# SPDX-FileCopyrightText: 2026 ZenteiQ AiTech Innovations Pvt. Ltd.
+# SPDX-License-Identifier: LGPL-3.0-only
+
+"""
+ScireX neural operators: a JAX/Flax library for learning solution
+operators of partial differential equations.
+
+Exposes the top-level model factory `get_model_jax` together with the FNO,
+WNO, GINO and Transolver architectures. Everything except the version string is
+imported lazily (PEP 562), so `import scirex` does not import JAX; the first
+access of `scirex.FNO` or `scirex.models` does.
+"""
+
+from __future__ import annotations
+
+import importlib
+from typing import TYPE_CHECKING, Any
+
+from scirex._internal.aliases import install as _install_aliases
+from scirex._internal.lazy import resolve_lazy_attr
+
+# Old import paths of modules that moved keep working (with a DeprecationWarning);
+# the table is scirex._internal.aliases.MOVED.
+_install_aliases()
+
+if TYPE_CHECKING:
+    __version__: str  # read from the package metadata on first access, see __getattr__
+    from .models import FNO, GINO, WNO, Transolver, get_model_jax
+
+__all__ = ["FNO", "GINO", "WNO", "Transolver", "get_model_jax"]
+
+_LAZY_EXPORTS: dict[str, str] = {
+    "FNO": "scirex.models",
+    "GINO": "scirex.models",
+    "Transolver": "scirex.models",
+    "WNO": "scirex.models",
+    "get_model_jax": "scirex.models",
+}
+
+# Subpackages reachable as attributes of `scirex` without an explicit import.
+_SUBMODULES = frozenset(
+    {
+        "config",
+        "data",
+        "experimental",
+        "layers",
+        "losses",
+        "models",
+        "training",
+        "utils",
+    }
+)
+
+
+def __getattr__(name: str) -> Any:
+    if name == "__version__":
+        # defined once, in pyproject.toml; looked up on first use so that
+        # `import scirex` does not scan the installed distributions
+        from scirex._internal.version import package_version
+
+        globals()["__version__"] = package_version()
+        return globals()["__version__"]
+    if name in _SUBMODULES:
+        return importlib.import_module(f"{__name__}.{name}")
+    return resolve_lazy_attr(__name__, name, _LAZY_EXPORTS, globals())
+
+
+def __dir__() -> list[str]:
+    return sorted({*globals(), *__all__, *_SUBMODULES, "__version__"})

scirex/_internal/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-FileCopyrightText: 2026 ZenteiQ AiTech Innovations Pvt. Ltd.
+# SPDX-License-Identifier: LGPL-3.0-only
+
+"""Private implementation helpers of ScireX. No stability guarantees.
+
+Nothing in this package is public API: modules, names and signatures may
+change or disappear in any release. Do not import from it in user code.
+"""

scirex/_internal/aliases.py

@@ -0,0 +1,144 @@
+# SPDX-FileCopyrightText: 2026 ZenteiQ AiTech Innovations Pvt. Ltd.
+# SPDX-License-Identifier: LGPL-3.0-only
+
+"""Old import paths of modules that moved, kept working.
+
+When a module moves, its old dotted path stays importable: importing it gives the *same
+module object* as the new path (so `isinstance` checks, monkeypatching and private
+names behave identically) and emits a `DeprecationWarning` that names the new path and
+the release that removes the old one. Submodules follow their package, so
+`scirex.distributed.comm` resolves to `scirex.experimental.distributed.comm`.
+
+The table `MOVED` is the single list of moves. It also feeds
+`scirex.models._legacy_refs`, which translates the same paths when a checkpoint is
+loaded. Add a line here in the same change that moves a module.
+"""
+
+from __future__ import annotations
+
+import importlib
+import importlib.abc
+import importlib.util
+import sys
+from collections.abc import Sequence
+from importlib.machinery import ModuleSpec
+from types import ModuleType
+
+from scirex._internal.deprecation import deprecation_message, warn_deprecated
+
+#: release in which the old paths were deprecated, and the earliest that may drop them
+SINCE = "1.0.0"
+REMOVAL = "1.2.0"
+
+#: old dotted module path -> where the module lives now
+MOVED: dict[str, str] = {
+    # not yet covered by the stability promise
+    "scirex.distributed": "scirex.experimental.distributed",
+    "scirex.geometry": "scirex.experimental.geometry",
+    "scirex.models.giwo": "scirex.experimental.giwo",
+    "scirex.training.patching": "scirex.experimental.patching",
+    "scirex.tuning": "scirex.experimental.tuning",
+    "scirex.losses.ad_differentiation": "scirex.experimental.losses.ad_differentiation",
+    "scirex.losses.ad_pde_terms": "scirex.experimental.losses.ad_pde_terms",
+    "scirex.losses.exact_ad_differentiation": "scirex.experimental.losses.exact_ad_differentiation",
+    "scirex.losses.patch_quadrature_ad": "scirex.experimental.losses.patch_quadrature_ad",
+    "scirex.losses.quadrature": "scirex.experimental.losses.quadrature",
+    "scirex.losses.domain_membership": "scirex.experimental.losses.domain_membership",
+    "scirex.losses.galerkin_t6": "scirex.experimental.losses.galerkin_t6",
+    "scirex.losses.darcy_hybrid_loss": "scirex.experimental.losses.darcy_hybrid_loss",
+    "scirex.config.distributed": "scirex.experimental.distributed.config",
+    # private helpers
+    "scirex.data.web_utils": "scirex._internal.web_utils",
+    "scirex.layers.common.einsum_utils": "scirex._internal.einsum_utils",
+    "scirex.layers.common.segment_csr": "scirex._internal.segment_csr",
+    "scirex.training.trainer": "scirex.experimental.training.trainer",
+    "scirex.training.training_state": "scirex.experimental.training.training_state",
+    "scirex.training.torch_setup": "scirex._internal.torch_setup",
+    "scirex.training.wno_train_utils": "scirex._internal.wno_train_utils",
+}
+
+
+def moved_message(old: str, new: str) -> str:
+    """The text of the `DeprecationWarning` for the import of `old`."""
+    return deprecation_message(old, since=SINCE, removal=REMOVAL, alternative=new)
+
+
+class _AliasLoader(importlib.abc.Loader):
+    """Hands out the module that lives at the new path instead of loading a copy."""
+
+    def __init__(self, old: str, new: str) -> None:
+        self.old, self.new = old, new
+        self._real_spec: ModuleSpec | None = None
+
+    def create_module(self, spec: ModuleSpec) -> ModuleType:
+        module = importlib.import_module(self.new)
+        self._real_spec = module.__spec__
+        # stacklevel 3: this frame, warn_deprecated's, then the importing line (the
+        # import machinery's own frames are skipped)
+        warn_deprecated(
+            self.old,
+            since=SINCE,
+            removal=REMOVAL,
+            alternative=self.new,
+            stacklevel=3,
+        )
+        return module
+
+    def exec_module(self, module: ModuleType) -> None:
+        # the import system stamped the alias spec on the shared module; put the real
+        # one back so `module.__spec__.name` still agrees with `module.__name__`
+        if self._real_spec is not None:
+            module.__spec__ = self._real_spec
+
+
+class AliasFinder(importlib.abc.MetaPathFinder):
+    """Meta-path finder for the keys of a move table (and their submodules)."""
+
+    def __init__(self, table: dict[str, str] | None = None) -> None:
+        self.table = MOVED if table is None else table
+
+    def target(self, fullname: str) -> str | None:
+        """The new dotted path for `fullname`, or `None` if it was not moved."""
+        for old, new in self.table.items():
+            if fullname == old:
+                return new
+            if fullname.startswith(old + "."):
+                return new + fullname[len(old) :]
+        return None
+
+    def find_spec(
+        self,
+        fullname: str,
+        path: Sequence[str] | None = None,
+        target: ModuleType | None = None,
+    ) -> ModuleSpec | None:
+        new = self.target(fullname)
+        if new is None:
+            return None
+        try:
+            real = importlib.util.find_spec(new)
+        except (ImportError, ValueError):
+            return None
+        if real is None:
+            return None
+        return importlib.util.spec_from_loader(
+            fullname,
+            _AliasLoader(fullname, new),
+            is_package=real.submodule_search_locations is not None,
+        )
+
+
+def install(table: dict[str, str] | None = None) -> AliasFinder:
+    """Put an `AliasFinder` first on `sys.meta_path` (once) and return it.
+
+    It must come first: the ordinary path finder would otherwise find the submodules of
+    a moved package inside its new directory and load them a second time under the old
+    name.
+    """
+    wanted = MOVED if table is None else table
+    for finder in sys.meta_path:
+        if isinstance(finder, AliasFinder) and finder.table is wanted:
+            return finder
+    finder = AliasFinder(wanted)
+    sys.meta_path.insert(0, finder)
+    return finder

scirex/_internal/deprecation.py

@@ -0,0 +1,92 @@
+# SPDX-FileCopyrightText: 2026 ZenteiQ AiTech Innovations Pvt. Ltd.
+# SPDX-License-Identifier: LGPL-3.0-only
+
+"""Deprecation helpers implementing the policy in `doc/dev_guide/api_policy.md`.
+
+A deprecated public name keeps working for at least two minor releases and emits
+a `DeprecationWarning` that says when it was deprecated, when it can be removed
+and what to use instead.
+"""
+
+from __future__ import annotations
+
+import functools
+import warnings
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def deprecation_message(
+    name: str, *, since: str, removal: str, alternative: str | None = None
+) -> str:
+    """The standard text of the `DeprecationWarning` for `name`."""
+    message = f"{name} is deprecated since ScireX {since} and will be removed in {removal} or later."
+    if alternative:
+        message += f" Use {alternative} instead."
+    return message
+
+
+def warn_deprecated(
+    name: str,
+    *,
+    since: str,
+    removal: str,
+    alternative: str | None = None,
+    stacklevel: int = 2,
+) -> None:
+    """Emit the standard `DeprecationWarning` for `name`.
+
+    Args:
+        name: Dotted name of the deprecated object, e.g. `"scirex.training.resolve_dtype"`.
+        since: Release in which the deprecation started.
+        removal: Earliest release in which the name may be removed.
+        alternative: What users should use instead, if anything.
+        stacklevel: Passed to `warnings.warn`; `2` points at the caller of
+            the function that calls this helper.
+    """
+    message = deprecation_message(
+        name, since=since, removal=removal, alternative=alternative
+    )
+    warnings.warn(message, DeprecationWarning, stacklevel=stacklevel)
+
+
+def deprecated(
+    *, since: str, removal: str, alternative: str | None = None
+) -> Callable[[F], F]:
+    """Decorate a function or class so that using it emits a `DeprecationWarning`."""
+
+    def decorator(obj: F) -> F:
+        name = f"{obj.__module__}.{obj.__qualname__}"
+        if isinstance(obj, type):
+            original_init = obj.__init__  # type: ignore[misc]
+
+            @functools.wraps(original_init)
+            def init(self: Any, *args: Any, **kwargs: Any) -> None:
+                warn_deprecated(
+                    name,
+                    since=since,
+                    removal=removal,
+                    alternative=alternative,
+                    stacklevel=3,
+                )
+                original_init(self, *args, **kwargs)
+
+            obj.__init__ = init  # type: ignore[assignment, misc]
+            return obj
+
+        @functools.wraps(obj)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            warn_deprecated(
+                name,
+                since=since,
+                removal=removal,
+                alternative=alternative,
+                stacklevel=3,
+            )
+            return obj(*args, **kwargs)
+
+        return wrapper  # type: ignore[return-value]
+
+    return decorator

scirex/_internal/einsum_utils.py

@@ -0,0 +1,108 @@
+# SPDX-FileCopyrightText: 2026 ZenteiQ AiTech Innovations Pvt. Ltd.
+# SPDX-License-Identifier: LGPL-3.0-only
+
+"""
+Einsum utilities for half-precision complex tensors
+(JAX analogues of torch `view_as_real` / `view_as_complex`).
+"""
+
+import jax.numpy as jnp
+import opt_einsum
+import tensorly as tl
+from tensorly.plugins import use_opt_einsum
+
+from .tensorly_backend import use_jax_backend
+
+use_jax_backend()
+use_opt_einsum("optimal")
+
+
+def _view_as_real(a: jnp.ndarray) -> jnp.ndarray:
+    """JAX equivalent of torch.view_as_real: complex (...) -> real (..., 2)"""
+    return jnp.stack([a.real, a.imag], axis=-1)
+
+
+def _view_as_complex(a: jnp.ndarray) -> jnp.ndarray:
+    """JAX equivalent of torch.view_as_complex: real (..., 2) -> complex (...)"""
+    return a[..., 0] + 1j * a[..., 1]
+
+
+def einsum_complexhalf_two_input(
+    eq: str, a: jnp.ndarray, b: jnp.ndarray
+) -> jnp.ndarray:
+    """
+    Compute (two-input) einsum for complexhalf tensors.
+    Because jnp.einsum currently does not support complex32 (complexhalf) types.
+    The inputs and outputs are the same as in jnp.einsum
+    """
+    assert len(eq.split(",")) == 2, "Equation must have two inputs."
+
+    # cast both tensors to "view as real" form, and half precision
+    a = _view_as_real(a)
+    b = _view_as_real(b)
+    a = a.astype(jnp.float16)
+    b = b.astype(jnp.float16)
+
+    # create a new einsum equation that takes into account "view as real" form
+    input_output = eq.split("->")
+    new_output = "xy" + input_output[1]
+    input_terms = input_output[0].split(",")
+    new_inputs = [input_terms[0] + "x", input_terms[1] + "y"]
+    new_eqn = new_inputs[0] + "," + new_inputs[1] + "->" + new_output
+
+    # convert back to complex form
+    tmp = tl.einsum(new_eqn, a, b)
+    res = jnp.stack(
+        [tmp[0, 0, ...] - tmp[1, 1, ...], tmp[1, 0, ...] + tmp[0, 1, ...]], axis=-1
+    )
+    return _view_as_complex(res)
+
+
+def einsum_complexhalf(eq: str, *args: jnp.ndarray) -> jnp.ndarray:
+    """
+    Compute einsum for complexhalf tensors.
+    Because jnp.einsum currently does not support complex32 (complexhalf) types.
+    The inputs and outputs are the same as in jnp.einsum
+    """
+    if len(args) == 2:
+        # if there are two inputs, it is faster to call this method
+        return einsum_complexhalf_two_input(eq, *args)
+
+    # find the optimal path
+    _, path_info = opt_einsum.contract_path(eq, *args)
+    partial_eqns = [
+        contraction_info[2] for contraction_info in path_info.contraction_list
+    ]
+
+    # create a dict of the input tensors by their label in the einsum equation
+    input_labels = eq.split("->")[0].split(",")
+    output_label = eq.split("->")[1]
+    tensors = dict(zip(input_labels, args))
+
+    # convert all tensors to half precision and "view as real" form
+    for key, tensor in tensors.items():
+        tensor = _view_as_real(tensor)
+        tensor = tensor.astype(jnp.float16)
+        tensors[key] = tensor
+
+    for partial_eq in partial_eqns:
+        # get the input tensors to partial_eq
+        in_labels, out_label = partial_eq.split("->")
+        in_labels = in_labels.split(",")
+        in_tensors = [tensors[label] for label in in_labels]
+
+        # create new einsum equation that takes into account "view as real" form
+        input_output = partial_eq.split("->")
+        new_output = "xy" + input_output[1]
+        input_terms = input_output[0].split(",")
+        new_inputs = [input_terms[0] + "x", input_terms[1] + "y"]
+        new_eqn = new_inputs[0] + "," + new_inputs[1] + "->" + new_output
+
+        # perform the einsum, and convert to "view as real" form
+        tmp = tl.einsum(new_eqn, *in_tensors)
+        result = jnp.stack(
+            [tmp[0, 0, ...] - tmp[1, 1, ...], tmp[1, 0, ...] + tmp[0, 1, ...]], axis=-1
+        )
+        tensors[out_label] = result
+
+    return _view_as_complex(tensors[output_label])

scirex/_internal/extras.py

@@ -0,0 +1,71 @@
+# SPDX-FileCopyrightText: 2026 ZenteiQ AiTech Innovations Pvt. Ltd.
+# SPDX-License-Identifier: LGPL-3.0-only
+
+"""Install hints for third-party packages that live in optional extras.
+
+`pyproject.toml` groups packages that only some features need (`vtk` for the
+benchmark loaders, `mlflow` and `optuna` for tuning) into extras. The helpers
+here turn the bare `ModuleNotFoundError` for such a package into an
+`ImportError` that names the feature and the extra to install.
+
+This module must stay free of heavy imports.
+"""
+
+from __future__ import annotations
+
+import functools
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def missing_extra_error(
+    err: ModuleNotFoundError, *, packages: tuple[str, ...], extra: str, feature: str
+) -> ImportError | None:
+    """Build an `ImportError` that names the extra to install.
+
+    Args:
+        err: The `ModuleNotFoundError` that was raised.
+        packages: Top-level import names that the extra provides.
+        extra: Name of the extra in `pyproject.toml`, e.g. `"hpo"`.
+        feature: Human-readable name of what needs the package.
+
+    Returns:
+        An `ImportError` with the install command, or `None` if the missing
+        module is not one of `packages` (so the original error should propagate).
+    """
+    missing = (err.name or "").split(".")[0]
+    if missing not in packages:
+        return None
+    return ImportError(
+        f"{feature} requires {missing!r}, which is not installed. "
+        f'Install the {extra!r} extra with `pip install "scirex[{extra}]"`.'
+    )
+
+
+def requires_extra(*packages: str, extra: str) -> Callable[[F], F]:
+    """Decorate a function so that a missing `packages` import names the extra.
+
+    The check happens when the function runs, not when it is defined, so code paths
+    that do not need the package (for example a cache hit) keep working without it.
+    """
+
+    def decorator(func: F) -> F:
+        feature = f"{func.__module__}.{func.__qualname__}"
+
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            try:
+                return func(*args, **kwargs)
+            except ModuleNotFoundError as err:
+                explained = missing_extra_error(
+                    err, packages=packages, extra=extra, feature=feature
+                )
+                if explained is None:
+                    raise
+                raise explained from err
+
+        return wrapper  # type: ignore[return-value]
+
+    return decorator

scirex/_internal/lazy.py

@@ -0,0 +1,109 @@
+# SPDX-FileCopyrightText: 2026 ZenteiQ AiTech Innovations Pvt. Ltd.
+# SPDX-License-Identifier: LGPL-3.0-only
+
+"""Helpers for package-level lazy exports (PEP 562).
+
+Public packages such as `scirex.models` export their names through a module
+level `__getattr__` so that `import scirex.models` does not import JAX, Flax or
+any optional dependency. The first access of a name imports the module that defines
+it. If that import fails because a third-party dependency is missing, the failure
+is reported as an `ImportError` that names the dependency and the feature
+that needs it.
+
+This module must stay free of heavy imports.
+"""
+
+from __future__ import annotations
+
+import importlib
+from collections.abc import Mapping, MutableMapping
+from typing import Any
+
+# Import name -> name on PyPI, for the packages whose absence we explain.
+_PIP_NAMES: dict[str, str] = {
+    "einops": "einops",
+    "jaxwt": "jaxwt",
+    "open3d": "open3d",
+    "scipy": "scipy",
+    "tensorly": "tensorly",
+}
+
+
+def missing_dependency_error(
+    err: ModuleNotFoundError, feature: str
+) -> ImportError | None:
+    """Build an explanatory `ImportError` for a missing third-party package.
+
+    Args:
+        err: The `ModuleNotFoundError` raised while importing `feature`.
+        feature: Human-readable name of what was being imported, e.g. `"WNO"`.
+
+    Returns:
+        An `ImportError` naming the missing package, or `None` if the missing
+        module is part of ScireX itself (a real bug that must not be disguised).
+    """
+    missing = (err.name or "").split(".")[0]
+    if not missing or missing == "scirex":
+        return None
+    pip_name = _PIP_NAMES.get(missing, missing)
+    return ImportError(
+        f"{feature} requires the third-party package {missing!r}, "
+        "which is not installed. "
+        f"Install it with `pip install {pip_name}`."
+    )
+
+
+def import_feature(module_name: str, feature: str) -> Any:
+    """Import `module_name`, explaining a missing third-party dependency.
+
+    Args:
+        module_name: Absolute dotted module name.
+        feature: Human-readable name of what needs the module.
+
+    Returns:
+        The imported module.
+
+    Raises:
+        ImportError: If a third-party dependency of the module is not installed.
+    """
+    try:
+        return importlib.import_module(module_name)
+    except ModuleNotFoundError as err:
+        explained = missing_dependency_error(err, feature)
+        if explained is None:
+            raise
+        raise explained from err
+
+
+def resolve_lazy_attr(
+    package: str,
+    name: str,
+    table: Mapping[str, str],
+    namespace: MutableMapping[str, Any],
+) -> Any:
+    """Resolve `name` of `package` through `table` and cache it in `namespace`.
+
+    Meant to be called from a module level `__getattr__`.
+
+    Args:
+        package: `__name__` of the package whose attribute is requested.
+        name: The requested attribute.
+        table: Maps each lazily exported name to the absolute name of the module
+            that provides it.
+        namespace: The package's `globals()`; the resolved value is stored there
+            so that later lookups bypass `__getattr__`.
+
+    Returns:
+        The attribute.
+
+    Raises:
+        AttributeError: If `name` is not in `table`.
+        ImportError: If a third-party dependency of the defining module is missing.
+    """
+    try:
+        module_name = table[name]
+    except KeyError:
+        raise AttributeError(f"module {package!r} has no attribute {name!r}") from None
+    value = getattr(import_feature(module_name, name), name)
+    namespace[name] = value
+    return value