spherical-deepkriging 1.0.4__tar.gz

spherical_deepkriging-1.0.4/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 STLABTW
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

spherical_deepkriging-1.0.4/MANIFEST.in

@@ -0,0 +1,5 @@
+prune spherical_deepkriging/basis_functions/mrts_sphere/cpp_extensions/build
+global-exclude *.so *.pyd *.dylib *.dll *.a *.lib
+
+# Include C++ sources for the spherical basis extension in sdist builds.
+recursive-include spherical_deepkriging/basis_functions/mrts_sphere/cpp_extensions *.cpp *.h *.hpp *.txt *.py CMakeLists.txt

spherical_deepkriging-1.0.4/PKG-INFO

@@ -0,0 +1,115 @@
+Metadata-Version: 2.4
+Name: spherical-deepkriging
+Version: 1.0.4
+Summary: Deep learning package for FRK
+Author-email: Wen-Ting Wang <egpivo@gmail.com>, Hao-Yun Huang <hhuscout@gms.ndhu.edu.tw>, Ping-Hsun Chiang <andrew501228@gmail.com>, Wu-Wei Ying <wuweiying1011@gms.ndhu.edu.tw>
+License-Expression: MIT
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch<3.0.0,>=2.0.0
+Requires-Dist: pytorch-lightning<3.0.0,>=2.6.0
+Requires-Dist: numpy<2.0.0,>=1.24.0
+Requires-Dist: tensorflow<3.0.0,>=2.19.0
+Requires-Dist: pandas>=2.0.0
+Requires-Dist: tabulate>=0.9.0
+Requires-Dist: scikit-learn>=1.3.0
+Requires-Dist: scipy>=1.10.0
+Requires-Dist: pybind11<3.0.0,>=2.11.0
+Requires-Dist: tqdm>=4.65.0
+Requires-Dist: tensorboard<2.20.0,>=2.13.0
+Requires-Dist: protobuf<5.0.0,>=4.25.8
+Requires-Dist: pykrige>=1.7.0
+Requires-Dist: seaborn>=0.13.0
+Requires-Dist: GPy>=1.10.0
+Requires-Dist: optuna>=3.0.0
+Requires-Dist: plotly>=5.0.0
+Requires-Dist: ipykernel<7.0.0,>=6.13.0
+Provides-Extra: vision
+Requires-Dist: torchvision>=0.15.0; extra == "vision"
+Provides-Extra: audio
+Requires-Dist: torchaudio>=2.0.0; extra == "audio"
+Provides-Extra: viz
+Requires-Dist: matplotlib>=3.7.0; extra == "viz"
+Requires-Dist: pillow<12.0.0,>=10.0.0; extra == "viz"
+Requires-Dist: plotly>=5.0.0; extra == "viz"
+Provides-Extra: data
+Provides-Extra: mrts
+Requires-Dist: jax<0.5.0,>=0.4.0; extra == "mrts"
+Requires-Dist: jaxlib<0.5.0,>=0.4.0; extra == "mrts"
+Provides-Extra: utils
+Provides-Extra: jupyter
+Requires-Dist: jupyter>=1.0.0; extra == "jupyter"
+Requires-Dist: ipykernel<7.0.0,>=6.13.0; extra == "jupyter"
+Requires-Dist: notebook>=7.0.0; extra == "jupyter"
+Requires-Dist: matplotlib>=3.7.0; extra == "jupyter"
+Provides-Extra: all
+Requires-Dist: torchvision>=0.15.0; extra == "all"
+Requires-Dist: torchaudio>=2.0.0; extra == "all"
+Requires-Dist: matplotlib>=3.7.0; extra == "all"
+Requires-Dist: scikit-learn>=1.3.0; extra == "all"
+Requires-Dist: scipy>=1.10.0; extra == "all"
+Requires-Dist: tqdm>=4.65.0; extra == "all"
+Requires-Dist: tensorboard<2.20.0,>=2.13.0; extra == "all"
+Requires-Dist: protobuf<5.0.0,>=4.25.8; extra == "all"
+Requires-Dist: pillow<12.0.0,>=10.0.0; extra == "all"
+Requires-Dist: pandas>=2.0.0; extra == "all"
+Requires-Dist: plotly>=5.0.0; extra == "all"
+Provides-Extra: dev
+Requires-Dist: twine>=4.0.0; extra == "dev"
+Requires-Dist: pytest>=7.4.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.1.0; extra == "dev"
+Requires-Dist: black>=23.7.0; extra == "dev"
+Requires-Dist: ruff>=0.1.0; extra == "dev"
+Requires-Dist: mypy>=1.5.0; extra == "dev"
+Dynamic: license-file
+Dynamic: requires-python
+
+# Spherical DeepKriging
+
+[![Tests](https://github.com/STLABTW/spherical-deepkriging/workflows/Test/badge.svg)](https://github.com/STLABTW/spherical-deepkriging/actions)
+[![codecov](https://codecov.io/github/STLABTW/spherical-deepkriging/graph/badge.svg?token=OF0LKVDII6)](https://codecov.io/github/STLABTW/spherical-deepkriging)
+
+Code for **DeepKriging on the Global Data** ([arXiv:2604.01689](https://arxiv.org/abs/2604.01689)): spherical spatial prediction with DeepKriging, MRTS-sphere / Wendland bases, and universal kriging. Implementation lives under `spherical_deepkriging/`.
+
+## Setup
+
+Needs [Miniconda](https://docs.conda.io/en/latest/miniconda.html). On Windows, use WSL.
+
+```bash
+git clone https://github.com/STLABTW/spherical-deepkriging.git
+cd spherical-deepkriging
+make install-dev
+```
+
+`make install-dev` creates the conda env and installs dependencies; `make build-spherical-cpp` builds the MRTS-sphere C++ extension.
+
+## Examples
+
+- Smoke test: `examples/toy/toy_sphere_deepkriging.ipynb`
+- Simulations: `examples/simulation/`
+- Real data: `examples/real_data/`
+
+See `examples/README.md` for run notes.
+
+## Citation
+
+```bibtex
+@misc{huang2026deepkrigingglobaldata,
+      title={DeepKriging on the Global Data},
+      author={Hao-Yun Huang and Wen-Ting Wang and Ping-Hsun Chiang and Wei-Ying Wu},
+      year={2026},
+      eprint={2604.01689},
+      archivePrefix={arXiv},
+      primaryClass={stat.ME},
+      url={https://arxiv.org/abs/2604.01689},
+}
+```

spherical_deepkriging-1.0.4/README.md

@@ -0,0 +1,40 @@
+# Spherical DeepKriging
+
+[![Tests](https://github.com/STLABTW/spherical-deepkriging/workflows/Test/badge.svg)](https://github.com/STLABTW/spherical-deepkriging/actions)
+[![codecov](https://codecov.io/github/STLABTW/spherical-deepkriging/graph/badge.svg?token=OF0LKVDII6)](https://codecov.io/github/STLABTW/spherical-deepkriging)
+
+Code for **DeepKriging on the Global Data** ([arXiv:2604.01689](https://arxiv.org/abs/2604.01689)): spherical spatial prediction with DeepKriging, MRTS-sphere / Wendland bases, and universal kriging. Implementation lives under `spherical_deepkriging/`.
+
+## Setup
+
+Needs [Miniconda](https://docs.conda.io/en/latest/miniconda.html). On Windows, use WSL.
+
+```bash
+git clone https://github.com/STLABTW/spherical-deepkriging.git
+cd spherical-deepkriging
+make install-dev
+```
+
+`make install-dev` creates the conda env and installs dependencies; `make build-spherical-cpp` builds the MRTS-sphere C++ extension.
+
+## Examples
+
+- Smoke test: `examples/toy/toy_sphere_deepkriging.ipynb`
+- Simulations: `examples/simulation/`
+- Real data: `examples/real_data/`
+
+See `examples/README.md` for run notes.
+
+## Citation
+
+```bibtex
+@misc{huang2026deepkrigingglobaldata,
+      title={DeepKriging on the Global Data},
+      author={Hao-Yun Huang and Wen-Ting Wang and Ping-Hsun Chiang and Wei-Ying Wu},
+      year={2026},
+      eprint={2604.01689},
+      archivePrefix={arXiv},
+      primaryClass={stat.ME},
+      url={https://arxiv.org/abs/2604.01689},
+}
+```

spherical_deepkriging-1.0.4/pyproject.toml

@@ -0,0 +1,145 @@
+[build-system]
+requires = [
+    "setuptools>=77.0.0",
+    "wheel",
+    "cmake>=3.10",
+    "pybind11>=2.11.0,<3.0.0",
+    "numpy>=1.24.0,<2.0.0",
+    "tomli>=2.0.0",
+]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "spherical-deepkriging"
+version = "1.0.4"
+description = "Deep learning package for FRK"
+readme = "README.md"
+requires-python = ">=3.10"
+license = "MIT"
+license-files = ["LICENSE"]
+authors = [
+    {name = "Wen-Ting Wang", email = "egpivo@gmail.com"},
+    {name = "Hao-Yun Huang", email = "hhuscout@gms.ndhu.edu.tw"},
+    {name = "Ping-Hsun Chiang", email = "andrew501228@gmail.com"},
+    {name = "Wu-Wei Ying", email = "wuweiying1011@gms.ndhu.edu.tw"},
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+
+dependencies = [
+    "torch>=2.0.0,<3.0.0",  # Avoid conflicts with neural-spatial (requires <2.3.0), but allow newer versions for DeepFRK
+    "pytorch-lightning>=2.6.0,<3.0.0",
+    "numpy>=1.24.0,<2.0.0",  # Avoid conflicts with packages requiring numpy<2.0.0 (hangman, neural-spatial, tensorflow)
+    "tensorflow>=2.19.0,<3.0.0",
+    "pandas>=2.0.0",
+    "tabulate>=0.9.0",
+    "scikit-learn>=1.3.0",
+    "scipy>=1.10.0",
+    "pybind11>=2.11.0,<3.0.0",  # Avoid conflicts with neural-spatial (requires <3.0.0)
+    "tqdm>=4.65.0",
+    "tensorboard>=2.13.0,<2.20.0",  # Use <2.20.0 to avoid conflicts, but keep >=2.13.0 for compatibility
+    "protobuf>=4.25.8,<5.0.0",
+    "pykrige>=1.7.0",  # Universal Kriging for spatial interpolation
+    "seaborn>=0.13.0",  # Statistical data visualization
+    "GPy>=1.10.0",  # Gaussian Processes for data generation (matching Lin paper)
+    "optuna>=3.0.0",  # Bayesian hyperparameter optimization
+    "plotly>=5.0.0",  # Required for optuna.visualization
+    "ipykernel>=6.13.0,<7.0.0",  # Jupyter kernel support
+]
+
+[project.optional-dependencies]
+vision = [
+    "torchvision>=0.15.0",
+]
+audio = [
+    "torchaudio>=2.0.0",
+]
+viz = [
+    "matplotlib>=3.7.0",
+    "pillow>=10.0.0,<12.0.0",  # Avoid conflicts with streamlit (requires <12)
+    "plotly>=5.0.0",  # Interactive visualization
+]
+data = [
+    # pandas and scikit-learn are now in main dependencies
+]
+mrts = [
+    "jax>=0.4.0,<0.5.0",  # Avoid conflicts with neural-spatial (requires <0.5.0)
+    "jaxlib>=0.4.0,<0.5.0",  # Avoid conflicts with neural-spatial (requires <0.5.0)
+]
+utils = [
+    # tqdm and tensorboard moved to main dependencies
+]
+jupyter = [
+    "jupyter>=1.0.0",
+    "ipykernel>=6.13.0,<7.0.0",  # Jupyter kernel support
+    "notebook>=7.0.0",
+    "matplotlib>=3.7.0",
+]
+all = [
+    "torchvision>=0.15.0",
+    "torchaudio>=2.0.0",
+    "matplotlib>=3.7.0",
+    "scikit-learn>=1.3.0",
+    "scipy>=1.10.0",
+    "tqdm>=4.65.0",
+    "tensorboard>=2.13.0,<2.20.0",
+    "protobuf>=4.25.8,<5.0.0",
+    "pillow>=10.0.0,<12.0.0",
+    "pandas>=2.0.0",
+    "plotly>=5.0.0",
+]
+dev = [
+    "twine>=4.0.0",
+    "pytest>=7.4.0",
+    "pytest-cov>=4.1.0",
+    "black>=23.7.0",
+    "ruff>=0.1.0",
+    "mypy>=1.5.0",
+]
+
+[tool.setuptools]
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["spherical_deepkriging*"]
+
+[tool.black]
+line-length = 88
+target-version = ['py310']
+
+[tool.ruff]
+line-length = 88
+target-version = "py310"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+python_files = "test_*.py"
+python_classes = "Test*"
+python_functions = "test_*"
+
+[tool.coverage.run]
+source = ["spherical_deepkriging"]
+omit = [
+    "*/tests/*",
+    "*/__pycache__/*",
+    "*/examples/*",
+]
+
+[tool.coverage.report]
+exclude_lines = [
+    "pragma: no cover",
+    "def __repr__",
+    "raise AssertionError",
+    "raise NotImplementedError",
+    "if __name__ == .__main__.:",
+    "if TYPE_CHECKING:",
+    "^\\.\\.\\.$",  # Exclude ellipsis in overload signatures
+]

spherical_deepkriging-1.0.4/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

spherical_deepkriging-1.0.4/setup.py

@@ -0,0 +1,112 @@
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+from setuptools import Extension, find_packages, setup
+from setuptools.command.build_ext import build_ext
+
+try:
+    import tomllib  # Python 3.11+
+except ModuleNotFoundError:  # pragma: no cover
+    import tomli as tomllib  # type: ignore[no-redef]
+
+
+ROOT = Path(__file__).resolve().parent
+
+
+class CMakeExtension(Extension):
+    def __init__(self, name: str, sourcedir: str):
+        super().__init__(name=name, sources=[])
+        self.sourcedir = str(Path(sourcedir).resolve())
+
+
+class CMakeBuild(build_ext):
+    def build_extension(self, ext: Extension) -> None:
+        if not isinstance(ext, CMakeExtension):
+            return super().build_extension(ext)
+
+        extdir = Path(self.get_ext_fullpath(ext.name)).parent.resolve()
+        build_temp = Path(self.build_temp) / ext.name.split(".")[-1]
+        build_temp.mkdir(parents=True, exist_ok=True)
+
+        cfg = "Release"
+        python_exe = sys.executable
+        jobs = max(1, (os.cpu_count() or 2) - 1)
+
+        cmake_args = [
+            f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}",
+            f"-DCMAKE_BUILD_TYPE={cfg}",
+            f"-DPython3_EXECUTABLE={python_exe}",
+        ]
+
+        # Ensure CMake can locate pybind11Config.cmake in isolated build envs.
+        try:
+            import pybind11  # type: ignore
+
+            pybind11_dir = Path(pybind11.get_cmake_dir()).resolve()
+            cmake_args.append(f"-Dpybind11_DIR={pybind11_dir}")
+        except Exception:
+            # Let CMake try its default discovery path if pybind11 is unavailable.
+            pass
+
+        subprocess.check_call(
+            ["cmake", "-S", ext.sourcedir, "-B", str(build_temp), *cmake_args]
+        )
+
+        # cmake --build passes `--config` only on multi-config generators (Windows)
+        if os.name == "nt":
+            subprocess.check_call(
+                ["cmake", "--build", str(build_temp), "--config", cfg]
+            )
+        else:
+            subprocess.check_call(
+                [
+                    "cmake",
+                    "--build",
+                    str(build_temp),
+                    "--config",
+                    cfg,
+                    "--",
+                    f"-j{jobs}",
+                ]
+            )
+
+
+def _load_project_metadata() -> dict:
+    with open(ROOT / "pyproject.toml", "rb") as f:
+        data = tomllib.load(f)
+    return data.get("project", {})  # type: ignore[no-any-return]
+
+
+project = _load_project_metadata()
+
+name = project.get("name", "spherical-deepkriging")
+version = project.get("version", "0.0.0")
+dependencies = project.get("dependencies", [])
+python_requires = project.get("requires-python", ">=3.10")
+
+cpp_ext_spherical_dir = (
+    ROOT
+    / "spherical_deepkriging"
+    / "basis_functions"
+    / "mrts_sphere"
+    / "cpp_extensions"
+)
+
+setup(
+    name=name,
+    version=version,
+    python_requires=python_requires,
+    install_requires=dependencies,
+    packages=find_packages(include=["spherical_deepkriging*"]),
+    include_package_data=True,
+    ext_modules=[
+        CMakeExtension(
+            "spherical_deepkriging.basis_functions.mrts_sphere.cpp_extensions.spherical_basis",
+            sourcedir=str(cpp_ext_spherical_dir),
+        )
+    ],
+    cmdclass={"build_ext": CMakeBuild},
+    zip_safe=False,
+)

spherical_deepkriging-1.0.4/spherical_deepkriging/basis_functions/mrts/mrts.py

@@ -0,0 +1,39 @@
+from typing import Dict, Optional, Tuple, Union
+
+import jax.numpy as jnp
+from jax.scipy.linalg import solve
+
+from spherical_deepkriging.basis_functions.mrts.utils import (
+    build_extended_matrix,
+    compute_h,
+    dist,
+    predict_rabf,
+)
+
+
+def mrts0(
+    knot: jnp.ndarray, k: int, x: Optional[jnp.ndarray] = None
+) -> Union[Tuple[jnp.ndarray, Dict[str, jnp.ndarray], int], jnp.ndarray]:
+    """Main function for MRTS with static pre-processing."""
+    Xu = jnp.asarray(knot)
+    ndims = Xu.shape[1]
+
+    if k < (ndims + 1):
+        raise ValueError(f"Invalid k: {k}. It must be >= {ndims + 1}")
+
+    slice_size = k - ndims - 1
+    AHA, UZ = build_extended_matrix(Xu, k=k, ndims=ndims, slice_size=slice_size)
+
+    BBBH = solve(
+        jnp.hstack([jnp.ones((Xu.shape[0], 1)), Xu]).T
+        @ jnp.hstack([jnp.ones((Xu.shape[0], 1)), Xu]),
+        jnp.hstack([jnp.ones((Xu.shape[0], 1)), Xu]).T,
+    ) @ compute_h(dist(Xu, Xu), ndims)
+    obj_attrs = {
+        "S": AHA,
+        "UZ": UZ,
+        "Xu": Xu,
+        "BBBH": BBBH,
+        "ndims": ndims,
+    }
+    return predict_rabf(obj_attrs, x, k) if x is not None else (Xu, obj_attrs, k)

spherical_deepkriging-1.0.4/spherical_deepkriging/basis_functions/mrts/utils.py

@@ -0,0 +1,127 @@
+from functools import partial
+from typing import Dict, Optional, Tuple
+
+import jax.numpy as jnp
+from jax import jit, lax
+
+
+@jit
+def dist(A: jnp.ndarray, B: jnp.ndarray) -> jnp.ndarray:
+    return jnp.sqrt(
+        jnp.sum((A[:, jnp.newaxis, :] - B[jnp.newaxis, :, :]) ** 2, axis=-1)
+    )
+
+
+def compute_h(d: jnp.ndarray, ndims: int) -> jnp.ndarray:
+    """Computes the H matrix based on the dimension."""
+
+    def case_ndims_1(_: None) -> jnp.ndarray:
+        return (1 / 12) * d**3
+
+    def case_ndims_2(_: None) -> jnp.ndarray:
+        return (1 / 8) * d**2 * jnp.log(d + 1e-8)
+
+    def case_ndims_3(_: None) -> jnp.ndarray:
+        return (-1 / 8) * d
+
+    cases = [case_ndims_1, case_ndims_2, case_ndims_3]
+    return lax.switch(ndims - 1, cases, None)
+
+
+@partial(jit, static_argnames=["ndims", "slice_size"])
+def build_extended_matrix(
+    Xu: jnp.ndarray, k: int, ndims: int, slice_size: int
+) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    """Builds the S and UZ matrices."""
+    n = Xu.shape[0]
+    B = jnp.hstack([jnp.ones((n, 1)), Xu])
+    BBB = jnp.linalg.solve(B.T @ B, B.T)
+    A = jnp.eye(n) - B @ BBB
+
+    d = jnp.sqrt(jnp.sum((Xu[:, None, :] - Xu[None, :, :]) ** 2, axis=-1))
+    H = compute_h(d, ndims)
+    AH = H - (H @ B) @ BBB
+    AHA = AH - BBB.T @ (B.T @ AH)
+
+    gamma0, _, _ = jnp.linalg.svd(AHA, full_matrices=False)
+
+    gamma0 = lax.dynamic_slice(
+        gamma0, start_indices=(0, 0), slice_sizes=(gamma0.shape[0], slice_size)
+    )
+
+    trueBS = AHA @ gamma0
+    rho = jnp.sqrt(jnp.sum(trueBS**2, axis=0))
+    gammas = A @ gamma0 / rho * jnp.sqrt(n)
+
+    extension_dim = ndims + 1
+
+    UZ = jnp.vstack(
+        [
+            jnp.hstack([gammas, jnp.zeros((gammas.shape[0], extension_dim))]),
+            jnp.zeros((extension_dim, gammas.shape[1] + extension_dim)),
+        ]
+    )
+    pad_size = ndims + 1
+
+    def compute_valid_indices() -> jnp.ndarray:
+        max_valid_index = min(pad_size, Xu.shape[1])
+        return jnp.arange(max_valid_index)
+
+    valid_indices = lax.stop_gradient(compute_valid_indices())
+
+    updates = 1 / jnp.std(Xu[:, valid_indices], axis=0) / jnp.sqrt((n - 1) / n)
+    UZ = UZ.at[n + valid_indices, k - pad_size - 1 + valid_indices].set(updates)
+
+    return AHA, UZ
+
+
+@partial(jit, static_argnames=["k"])
+def predict_rabf(
+    obj: Dict[str, jnp.ndarray],
+    newx: Optional[jnp.ndarray] = None,
+    k: Optional[int] = None,
+) -> jnp.ndarray:
+    """Predicts new values based on the provided model object."""
+    if newx is None:
+        return obj
+
+    x0 = newx
+    d = dist(x0, obj["Xu"])
+    ndims = x0.shape[1] if x0.ndim > 1 else 1
+    H = compute_h(d, ndims)
+
+    kstar = k - ndims - 1
+
+    def true_branch(_: None) -> jnp.ndarray:
+        slice_UZ = lax.dynamic_slice(
+            obj["UZ"], start_indices=(0, 0), slice_sizes=(obj["Xu"].shape[0], kstar)
+        )
+        X1 = H @ slice_UZ
+
+        B = jnp.hstack(
+            [jnp.ones((x0.shape[0], 1)), x0 if x0.ndim > 1 else x0[:, jnp.newaxis]]
+        )
+        BBBH_UZ = lax.dynamic_slice(
+            obj["UZ"], start_indices=(0, 0), slice_sizes=(obj["Xu"].shape[0], kstar)
+        )
+        X1 -= B @ obj["BBBH"] @ BBBH_UZ
+        X1 /= jnp.sqrt(obj["Xu"].shape[0])
+        return X1
+
+    def false_branch(_: None) -> jnp.ndarray:
+        return jnp.zeros((x0.shape[0], kstar))
+
+    X1 = lax.cond(kstar > 0, true_branch, false_branch, operand=None)
+
+    X2 = jnp.hstack(
+        [jnp.ones((x0.shape[0], 1)), x0 if x0.ndim > 1 else x0[:, jnp.newaxis]]
+    )
+
+    def concat_branch(_: None) -> jnp.ndarray:
+        return jnp.hstack([X2, X1])
+
+    def no_concat_branch(_: None) -> jnp.ndarray:
+        padding = jnp.zeros((x0.shape[0], X1.shape[1]))
+        return jnp.hstack([X2, padding])
+
+    return lax.cond(kstar > 0, concat_branch, no_concat_branch, operand=None)