jaccpot 0.0.1__py3-none-any.whl

jaccpot/__init__.py

@@ -0,0 +1,34 @@
+"""Jaccpot: high-level FMM solver APIs built on Yggdrasil artifacts."""
+
+from ._typecheck import enable_runtime_typecheck
+
+enable_runtime_typecheck()
+
+from .autodiff import differentiable_gravitational_acceleration
+from .basis import ComplexSHBasis, RealSHBasis
+from .config import (
+    FarFieldConfig,
+    FMMAdvancedConfig,
+    FMMPreset,
+    MemoryObjective,
+    NearFieldConfig,
+    RuntimePolicyConfig,
+    TreeConfig,
+)
+from .odisseo import OdisseoFMMCoupler
+from .solver import FastMultipoleMethod
+
+__all__ = [
+    "FMMAdvancedConfig",
+    "FMMPreset",
+    "FarFieldConfig",
+    "FastMultipoleMethod",
+    "ComplexSHBasis",
+    "MemoryObjective",
+    "RealSHBasis",
+    "NearFieldConfig",
+    "OdisseoFMMCoupler",
+    "RuntimePolicyConfig",
+    "TreeConfig",
+    "differentiable_gravitational_acceleration",
+]

jaccpot/_typecheck.py

@@ -0,0 +1,33 @@
+"""Runtime type-check configuration for jaccpot."""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+_TYPECHECK_HOOK: Any = None
+
+
+def _runtime_typecheck_enabled() -> bool:
+    raw = os.getenv("JACCPOT_RUNTIME_TYPECHECK", "0").strip().lower()
+    return raw not in {"0", "false", "no", "off"}
+
+
+def enable_runtime_typecheck() -> bool:
+    """Enable package-wide jaxtyping+beartype checks for annotated callables."""
+    global _TYPECHECK_HOOK
+
+    if _TYPECHECK_HOOK is not None:
+        return True
+    if not _runtime_typecheck_enabled():
+        return False
+
+    from jaxtyping import install_import_hook
+
+    # Instruments submodule imports under `jaccpot` so annotated callables are
+    # checked by beartype with jaxtyping's shape/dtype semantics.
+    _TYPECHECK_HOOK = install_import_hook("jaccpot", typechecker="beartype.beartype")
+    return True
+
+
+__all__ = ["enable_runtime_typecheck"]

jaccpot/autodiff.py

@@ -0,0 +1,36 @@
+"""Autodiff helpers for Jaccpot."""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import jax
+import jax.numpy as jnp
+from jaxtyping import Array
+
+
+def differentiable_gravitational_acceleration(
+    positions: Array,
+    masses: Array,
+    *,
+    theta: float = 0.6,
+    G: float = 1.0,
+    softening: float = 1e-3,
+    bounds: Optional[Tuple[Array, Array]] = None,
+    leaf_size: int = 16,
+    max_order: int = 4,
+) -> Array:
+    """Differentiable gravitational accelerations via direct summation."""
+
+    del theta, bounds, leaf_size, max_order
+
+    diffs = positions[:, None, :] - positions[None, :, :]
+    dist2 = jnp.sum(diffs * diffs, axis=-1) + softening**2
+    inv_dist = jnp.where(dist2 > 0, jnp.power(dist2, -0.5), 0.0)
+    inv_dist3 = inv_dist**3
+    weights = masses[None, :] * inv_dist3
+    weights = weights * (1.0 - jnp.eye(positions.shape[0], dtype=positions.dtype))
+    return -G * jnp.einsum("ij,ijk->ik", weights, diffs)
+
+
+__all__ = ["differentiable_gravitational_acceleration"]

jaccpot/basis/__init__.py

@@ -0,0 +1,14 @@
+"""Basis abstractions for FMM coefficient representations."""
+
+from .base import BasisInterface, BasisMetadata
+from .complex_sh import ComplexSHBasis
+from .real_sh import RealSHBasis, complex_to_real_coeffs, real_to_complex_coeffs
+
+__all__ = [
+    "BasisInterface",
+    "BasisMetadata",
+    "ComplexSHBasis",
+    "RealSHBasis",
+    "complex_to_real_coeffs",
+    "real_to_complex_coeffs",
+]

jaccpot/basis/base.py

@@ -0,0 +1,56 @@
+"""Shared basis interface for FMM coefficient transforms and M2L kernels."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Protocol, runtime_checkable
+
+from jaxtyping import Array
+
+
+@runtime_checkable
+class BasisInterface(Protocol):
+    """Minimal basis contract used by solver/runtime orchestration.
+
+    Implementations define coefficient layout metadata and batched transforms
+    needed by M2L operators.
+    """
+
+    name: str
+    p_max: int
+    coefficient_ordering: str
+    runtime_expansion_basis: str
+
+    def n_coeffs(self: "BasisInterface", p: int) -> int:
+        """Return number of packed coefficients for expansion order ``p``."""
+
+    def pack_coeffs(self: "BasisInterface", coeffs: Array, *, order: int) -> Array:
+        """Pack basis coefficients into the runtime 1D layout."""
+
+    def unpack_coeffs(self: "BasisInterface", packed: Array, *, order: int) -> Array:
+        """Unpack runtime 1D coefficient layout into structured basis form."""
+
+    def rotate_to_z(
+        self: "BasisInterface", coeffs: Array, directions: Array, *, order: int
+    ) -> Array:
+        """Rotate batched coefficients into a frame where direction maps to ``+z``."""
+
+    def rotate_from_z(
+        self: "BasisInterface", coeffs: Array, directions: Array, *, order: int
+    ) -> Array:
+        """Rotate batched coefficients back from the ``+z`` frame."""
+
+    def m2l_rot_scale(
+        self: "BasisInterface", sources: Array, deltas: Array, *, order: int
+    ) -> Array:
+        """Evaluate batched M2L translation in the basis' native convention."""
+
+
+@dataclass(frozen=True)
+class BasisMetadata:
+    """Static metadata helper for basis implementations."""
+
+    name: str
+    p_max: int
+    coefficient_ordering: str
+    runtime_expansion_basis: str

jaccpot/basis/complex_sh.py

@@ -0,0 +1,114 @@
+"""Complex spherical-harmonic basis adapter for existing solidfmm kernels."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import jax
+import jax.numpy as jnp
+from jaxtyping import Array
+
+from jaccpot.operators.complex_ops import (
+    complex_rotation_blocks_from_z_solidfmm_batch,
+    complex_rotation_blocks_to_z_solidfmm_batch,
+    m2l_complex_reference_batch,
+)
+from jaccpot.operators.real_harmonics import sh_size
+
+
+@dataclass(frozen=True)
+class ComplexSHBasis:
+    """Adapter exposing the current solidfmm complex basis through a common API.
+
+    Coefficients use the existing packed solidfmm layout (ell-major blocks with
+    ``m=-ell..+ell`` for each degree).
+    """
+
+    p_max: int = 32
+    name: str = "complex"
+    coefficient_ordering: str = "ell-major, m=-ell..+ell (packed complex SH)"
+    runtime_expansion_basis: str = "solidfmm"
+
+    def n_coeffs(self: "ComplexSHBasis", p: int) -> int:
+        """Return packed coefficient count for order ``p``."""
+        return int(sh_size(int(p)))
+
+    def pack_coeffs(self: "ComplexSHBasis", coeffs: Array, *, order: int) -> Array:
+        """Validate and return packed solidfmm complex coefficients."""
+        coeffs_arr = jnp.asarray(coeffs)
+        expected = self.n_coeffs(order)
+        if int(coeffs_arr.shape[-1]) != expected:
+            raise ValueError(
+                f"expected last dimension {expected} for order={int(order)}, got {coeffs_arr.shape[-1]}"
+            )
+        return coeffs_arr
+
+    def unpack_coeffs(self: "ComplexSHBasis", packed: Array, *, order: int) -> Array:
+        """Return packed coefficients (complex basis already uses packed layout)."""
+        return self.pack_coeffs(packed, order=order)
+
+    def rotate_to_z(
+        self: "ComplexSHBasis", coeffs: Array, directions: Array, *, order: int
+    ) -> Array:
+        """Rotate batched multipoles into a frame aligned to ``+z``."""
+        coeffs_arr = self.pack_coeffs(coeffs, order=order)
+        x, y, z = self._split_xyz(directions)
+        to_blocks = complex_rotation_blocks_to_z_solidfmm_batch(
+            x,
+            y,
+            z,
+            order=int(order),
+            dtype=coeffs_arr.dtype,
+        )
+        return self._apply_rotation_blocks(coeffs_arr, to_blocks)
+
+    def rotate_from_z(
+        self: "ComplexSHBasis", coeffs: Array, directions: Array, *, order: int
+    ) -> Array:
+        """Rotate batched multipoles from the ``+z`` frame back to world frame."""
+        coeffs_arr = self.pack_coeffs(coeffs, order=order)
+        x, y, z = self._split_xyz(directions)
+        from_blocks = complex_rotation_blocks_from_z_solidfmm_batch(
+            x,
+            y,
+            z,
+            order=int(order),
+            dtype=coeffs_arr.dtype,
+        )
+        return self._apply_rotation_blocks(coeffs_arr, from_blocks)
+
+    def m2l_rot_scale(
+        self: "ComplexSHBasis", sources: Array, deltas: Array, *, order: int
+    ) -> Array:
+        """Run existing batched complex M2L kernel for source multipoles."""
+        src = self.pack_coeffs(sources, order=order)
+        delta_arr = jnp.asarray(deltas)
+        if delta_arr.ndim != 2 or int(delta_arr.shape[1]) != 3:
+            raise ValueError("deltas must have shape (batch, 3)")
+        return m2l_complex_reference_batch(src, delta_arr, order=int(order))
+
+    @staticmethod
+    def _split_xyz(directions: Array) -> tuple[Array, Array, Array]:
+        directions_arr = jnp.asarray(directions)
+        if directions_arr.ndim != 2 or int(directions_arr.shape[1]) != 3:
+            raise ValueError("directions must have shape (batch, 3)")
+        return directions_arr[:, 0], directions_arr[:, 1], directions_arr[:, 2]
+
+    @staticmethod
+    def _apply_rotation_blocks(coeffs: Array, blocks: list[Array]) -> Array:
+        offsets = [0]
+        for ell, block in enumerate(blocks):
+            width = int(block.shape[-1])
+            if width != 2 * ell + 1:
+                raise ValueError("rotation block width does not match ell")
+            offsets.append(offsets[-1] + width)
+
+        def rotate_one(row: Array, row_blocks: list[Array]) -> Array:
+            parts = []
+            for ell, block in enumerate(row_blocks):
+                start = offsets[ell]
+                end = offsets[ell + 1]
+                parts.append(block @ row[start:end])
+            return jnp.concatenate(parts, axis=0)
+
+        return jax.vmap(rotate_one)(coeffs, blocks)

jaccpot/basis/real_sh.py

@@ -0,0 +1,142 @@
+"""Real spherical-harmonic basis layout and complex/real conversion helpers."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import jax.numpy as jnp
+from jaxtyping import Array
+
+
+def _idx_nm(n: int, m: int) -> int:
+    """Packed index for degree ``n`` and order ``m`` in ``m=-n..n`` layout."""
+    if abs(int(m)) > int(n):
+        raise ValueError("|m| must be <= n")
+    return int(n) * int(n) + (int(m) + int(n))
+
+
+def n_real_sh_coeffs(order: int) -> int:
+    """Number of packed real SH coefficients for ``order``."""
+    p = int(order)
+    if p < 0:
+        raise ValueError("order must be non-negative")
+    return (p + 1) * (p + 1)
+
+
+def complex_to_real_coeffs(complex_coeffs: Array, *, order: int) -> Array:
+    """Convert packed complex SH coefficients to packed real SH coefficients.
+
+    This conversion assumes Condon-Shortley style paired complex coefficients
+    where each ``(n, m>0)`` pair is transformed into two real coefficients
+    ``(n, +m)`` and ``(n, -m)``.
+    """
+
+    coeffs = jnp.asarray(complex_coeffs)
+    expected = n_real_sh_coeffs(order)
+    if int(coeffs.shape[-1]) != expected:
+        raise ValueError(
+            f"expected last dimension {expected} for order={int(order)}, got {coeffs.shape[-1]}"
+        )
+
+    out = jnp.zeros(coeffs.shape, dtype=coeffs.real.dtype)
+    sqrt2 = jnp.sqrt(jnp.asarray(2.0, dtype=coeffs.real.dtype))
+
+    for n in range(int(order) + 1):
+        idx0 = _idx_nm(n, 0)
+        out = out.at[..., idx0].set(jnp.real(coeffs[..., idx0]))
+        for m in range(1, n + 1):
+            idx_p = _idx_nm(n, m)
+            idx_n = _idx_nm(n, -m)
+            sign = -1.0 if (m % 2) else 1.0
+            c_p = coeffs[..., idx_p]
+            c_n = coeffs[..., idx_n]
+            r_p = (sign * c_p + c_n) / sqrt2
+            r_n = (sign * c_p - c_n) / (1j * sqrt2)
+            out = out.at[..., idx_p].set(jnp.real(r_p))
+            out = out.at[..., idx_n].set(jnp.real(r_n))
+
+    return out
+
+
+def real_to_complex_coeffs(real_coeffs: Array, *, order: int) -> Array:
+    """Convert packed real SH coefficients to packed complex SH coefficients."""
+
+    coeffs = jnp.asarray(real_coeffs)
+    expected = n_real_sh_coeffs(order)
+    if int(coeffs.shape[-1]) != expected:
+        raise ValueError(
+            f"expected last dimension {expected} for order={int(order)}, got {coeffs.shape[-1]}"
+        )
+
+    out = jnp.zeros(coeffs.shape, dtype=jnp.result_type(coeffs.dtype, jnp.complex64))
+    sqrt2 = jnp.sqrt(jnp.asarray(2.0, dtype=coeffs.dtype))
+
+    for n in range(int(order) + 1):
+        idx0 = _idx_nm(n, 0)
+        out = out.at[..., idx0].set(coeffs[..., idx0].astype(out.dtype))
+        for m in range(1, n + 1):
+            idx_p = _idx_nm(n, m)
+            idx_n = _idx_nm(n, -m)
+            sign = -1.0 if (m % 2) else 1.0
+            r_p = coeffs[..., idx_p]
+            r_n = coeffs[..., idx_n]
+            c_p = sign * (r_p + 1j * r_n) / sqrt2
+            c_n = (r_p - 1j * r_n) / sqrt2
+            out = out.at[..., idx_p].set(c_p.astype(out.dtype))
+            out = out.at[..., idx_n].set(c_n.astype(out.dtype))
+
+    return out
+
+
+@dataclass(frozen=True)
+class RealSHBasis:
+    """Packed real spherical-harmonic basis metadata and layout helpers."""
+
+    p_max: int = 32
+    name: str = "real"
+    coefficient_ordering: str = "ell-major, m=-ell..+ell (packed real SH)"
+    runtime_expansion_basis: str = "solidfmm"
+
+    def n_coeffs(self: "RealSHBasis", p: int) -> int:
+        """Return packed coefficient count for expansion order ``p``."""
+        return n_real_sh_coeffs(int(p))
+
+    def pack_coeffs(self: "RealSHBasis", coeffs: Array, *, order: int) -> Array:
+        """Validate and return packed real SH coefficients."""
+        arr = jnp.asarray(coeffs)
+        expected = self.n_coeffs(order)
+        if int(arr.shape[-1]) != expected:
+            raise ValueError(
+                f"expected last dimension {expected} for order={int(order)}, got {arr.shape[-1]}"
+            )
+        return arr
+
+    def unpack_coeffs(self: "RealSHBasis", packed: Array, *, order: int) -> Array:
+        """Return packed coefficients (real SH uses the packed runtime layout)."""
+        return self.pack_coeffs(packed, order=order)
+
+    def rotate_to_z(
+        self: "RealSHBasis", coeffs: Array, directions: Array, *, order: int
+    ) -> Array:
+        """Rotate real SH coefficients into a z-aligned frame (not yet implemented)."""
+        raise NotImplementedError("real SH rotations are implemented in Stage J3")
+
+    def rotate_from_z(
+        self: "RealSHBasis", coeffs: Array, directions: Array, *, order: int
+    ) -> Array:
+        """Rotate real SH coefficients back from a z-aligned frame."""
+        raise NotImplementedError("real SH rotations are implemented in Stage J3")
+
+    def m2l_rot_scale(
+        self: "RealSHBasis", sources: Array, deltas: Array, *, order: int
+    ) -> Array:
+        """Translate real SH multipoles to locals via rotate/scale M2L path."""
+        raise NotImplementedError("real SH rotate+scale M2L is implemented in Stage J3")
+
+
+__all__ = [
+    "RealSHBasis",
+    "complex_to_real_coeffs",
+    "real_to_complex_coeffs",
+    "n_real_sh_coeffs",
+]

jaccpot/config.py

@@ -0,0 +1,115 @@
+"""Preset-first configuration model for Jaccpot."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Literal, Optional
+
+BASIS_DOC = "Preferred production basis is 'solidfmm'."
+FARFIELD_MODE_DOC = (
+    "For large_n_gpu production, far-field execution is canonicalized to "
+    "'pair_grouped'."
+)
+NEARFIELD_MODE_DOC = (
+    "For large_n_gpu production, near-field execution is canonicalized to "
+    "'bucketed'."
+)
+MEMORY_OBJECTIVE_DOC = (
+    "For large_n_gpu production, memory objective is canonicalized to "
+    "'minimum_memory'."
+)
+
+Basis = Literal["cartesian", "solidfmm", "complex", "real"]
+FarFieldMode = Literal["auto", "pair_grouped", "class_major"]
+NearFieldMode = Literal["auto", "baseline", "bucketed"]
+MemoryObjective = Literal["balanced", "throughput", "minimum_memory"]
+FMMExecutionBackend = Literal["auto", "radix", "octree"]
+
+
+class FMMPreset(str, Enum):
+    """User-facing quality/speed presets."""
+
+    FAST = "fast"
+    BALANCED = "balanced"
+    ACCURATE = "accurate"
+    LARGE_N_GPU = "large_n_gpu"
+
+
+@dataclass(frozen=True)
+class TreeConfig:
+    """Tree-construction overrides for advanced runtime tuning."""
+
+    tree_type: Optional[str] = None
+    mode: Optional[str] = None
+    leaf_target: Optional[int] = None
+    refine_local: Optional[bool] = None
+    max_refine_levels: Optional[int] = None
+    aspect_threshold: Optional[float] = None
+
+
+@dataclass(frozen=True)
+class FarFieldConfig:
+    """Far-field interaction and translation-kernel overrides."""
+
+    grouped_interactions: Optional[bool] = None
+    mode: FarFieldMode = "auto"
+    rotation: Optional[str] = None
+    m2l_chunk_size: Optional[int] = None
+    l2l_chunk_size: Optional[int] = None
+    streamed_far_pairs: Optional[bool] = None
+    mixed_order: bool = False
+    mixed_order_min_order: Optional[int] = None
+
+
+@dataclass(frozen=True)
+class NearFieldConfig:
+    """Near-field direct-interaction strategy overrides."""
+
+    mode: NearFieldMode = "auto"
+    edge_chunk_size: int = 256
+    precompute_scatter_schedules: bool = True
+
+
+@dataclass(frozen=True)
+class RuntimePolicyConfig:
+    """Execution-policy overrides for tree build and traversal.
+
+    Notes:
+    - `runtime_path='legacy'` is deprecated and will be removed.
+    - For `preset='large_n_gpu'`, runtime policy is canonicalized to the
+      production low-memory fast path (minimum_memory + streamed pair_grouped
+      + bucketed nearfield).
+    """
+
+    execution_backend: FMMExecutionBackend = "auto"
+    host_refine_mode: str = "auto"
+    fail_fast: bool = False
+    jit_tree: Optional[bool] = None
+    jit_traversal: Optional[bool] = None
+    memory_objective: MemoryObjective = "balanced"
+    memory_budget_bytes: Optional[int] = None
+    max_pair_queue: Optional[int] = None
+    pair_process_block: Optional[int] = None
+    traversal_config: Optional[Any] = None
+    enable_interaction_cache: bool = True
+    retain_traversal_result: bool = True
+    retain_interactions: bool = True
+    prepare_stage_memory_split_enabled: Optional[bool] = None
+    autotune_m2l_chunk: bool = False
+    precompute_grouped_class_segments: Optional[bool] = None
+    grouped_schedule_budget_bytes: Optional[int] = None
+    nearfield_schedule_item_cap: Optional[int] = None
+    upward_leaf_batch_size: Optional[int] = None
+
+
+@dataclass(frozen=True)
+class FMMAdvancedConfig:
+    """Aggregate container for all advanced FMM override groups."""
+
+    tree: TreeConfig = TreeConfig()
+    farfield: FarFieldConfig = FarFieldConfig()
+    nearfield: NearFieldConfig = NearFieldConfig()
+    runtime: RuntimePolicyConfig = RuntimePolicyConfig()
+    mac_type: Optional[str] = None
+    dehnen_radius_scale: float = 1.0

jaccpot/downward/__init__.py

@@ -0,0 +1 @@
+"""Downward-sweep operators for local-expansion propagation."""