dendra 0.27.2__py3-none-any.whl

dendra/__init__.py

@@ -0,0 +1,257 @@
+"""Dendra: A PyTorch-based framework for simulating and modeling neural dynamics."""
+
+import os
+
+__version__ = "0.27.2"
+
+os.environ.setdefault("OMP_DYNAMIC", "FALSE")
+os.environ.setdefault("MKL_DYNAMIC", "FALSE")
+os.environ["OMP_PROC_BIND"] = "true"
+os.environ["OMP_PLACES"] = "cores"
+
+from ._bootstrap import (
+    _should_cache_cpu_isa_for_dendra,
+    _should_configure_torchinductor_cache_on_import,
+    configure_torchinductor_cache_for_dendra,
+)
+
+_torchinductor_cache_opt_in = _should_configure_torchinductor_cache_on_import()
+if _torchinductor_cache_opt_in:
+    configure_torchinductor_cache_for_dendra()
+
+import pickle
+import time
+import warnings
+from pathlib import Path
+
+import torch
+
+import dendra.func as func
+import dendra.models.callbacks as callbacks
+import dendra.models.mod as mod
+
+from .const import FARADAY, PI, TAU, E, R
+from .diagnostics import doctor
+from .helpers import (
+    COMPILE_OPTIONS,
+    DEVICE,
+    DTYPE,
+    NATIVE_EXTENSION_POLICY,
+    PRESERVE_MECHANISM_POPULATION_AXIS,
+    RUNTIME_CONTRACT_VALIDATION,
+    TF32,
+    allow_tf32,
+    ctx,
+    current_compile_options,
+    current_device,
+    current_dtype,
+    current_native_extension_policy,
+    current_preserve_mechanism_population_axis,
+    current_runtime_contract_validation,
+    normalize_native_extension_policy,
+    normalize_preserve_mechanism_population_axis,
+    normalize_runtime_contract_validation,
+    set_compile_options,
+    set_jit_enabled,
+    set_jit_in_network_enabled,
+    set_jit_network_ops_enabled,
+    set_jit_network_solves_enabled,
+)
+from .models import (
+    Axon,
+    Cable,
+    CompartmentGeometry,
+    CompartmentGraph,
+    CompartmentMetadata,
+    CompartmentTopology,
+    DistributionSpec,
+    ExtCellAxon,
+    ExtCellTree,
+    Morphology,
+    Myelinated,
+    NetStim,
+    Network,
+    Population,
+    RandomParameterSpec,
+    RuntimeNoiseSpec,
+    Section,
+    SectionLocation,
+    SingleCompartment,
+    SynapseSlots,
+    Tree,
+    Unmyelinated,
+    available_random_distributions,
+    concat_models,
+    concat_slices,
+    connect_morphologies,
+    get_random_distribution,
+    insert_intrinsic_activity,
+    make_random_parameter_spec,
+    make_runtime_noise_spec,
+    passive_end_nodes_,
+    register_random_distribution,
+    remove_intrinsic_activity,
+    sample_random_parameter,
+    sample_runtime_noise,
+)
+from .models.fields import *
+from .models.integrators import (
+    bwd_euler_bt,
+    bwd_euler_sc,
+    bwd_euler_ub,
+    df,
+    dfh,
+    dhs,
+    dufort_frankel,
+    dufort_frankel_homogeneous,
+    euler,
+    eulerv1,
+    rk1,
+    rk2,
+    rk4,
+    scnv,
+)
+from .models.mod import load_mechanisms
+from .models.stim.waveform import *
+from .stepping import step, step_network, step_population
+
+__all__ = [
+    "PI",
+    "E",
+    "TAU",
+    "R",
+    "FARADAY",
+    "ctx",
+    "doctor",
+    "step",
+    "step_population",
+    "step_network",
+    "DEVICE",
+    "DTYPE",
+    "NATIVE_EXTENSION_POLICY",
+    "PRESERVE_MECHANISM_POPULATION_AXIS",
+    "RUNTIME_CONTRACT_VALIDATION",
+    "COMPILE_OPTIONS",
+    "current_compile_options",
+    "current_device",
+    "current_dtype",
+    "current_native_extension_policy",
+    "current_preserve_mechanism_population_axis",
+    "current_runtime_contract_validation",
+    "normalize_runtime_contract_validation",
+    "normalize_native_extension_policy",
+    "normalize_preserve_mechanism_population_axis",
+    "set_compile_options",
+    "set_jit_enabled",
+    "set_jit_network_solves_enabled",
+    "set_jit_network_ops_enabled",
+    "set_jit_in_network_enabled",
+    "callbacks",
+    "func",
+    "mod",
+    "load_mechanisms",
+    "DistributionSpec",
+    "RandomParameterSpec",
+    "RuntimeNoiseSpec",
+    "available_random_distributions",
+    "get_random_distribution",
+    "register_random_distribution",
+    "make_random_parameter_spec",
+    "make_runtime_noise_spec",
+    "sample_random_parameter",
+    "sample_runtime_noise",
+    "insert_intrinsic_activity",
+    "remove_intrinsic_activity",
+    "passive_end_nodes_",
+    "Population",
+    "Morphology",
+    "Section",
+    "SectionLocation",
+    "CompartmentTopology",
+    "CompartmentGeometry",
+    "CompartmentMetadata",
+    "CompartmentGraph",
+    "SingleCompartment",
+    "Cable",
+    "Axon",
+    "Unmyelinated",
+    "Myelinated",
+    "Tree",
+    "ExtCellAxon",
+    "ExtCellTree",
+    "NetStim",
+    "Network",
+    "SynapseSlots",
+    "concat_models",
+    "concat_slices",
+    "connect_morphologies",
+    "anisotropic_point",
+    "isotropic_point",
+    "parametric_efield",
+    "PreComputedInterpolate1D",
+    "precomputed_interpolate_1d",
+    "EfieldInterpolate3DRect",
+    "efield_interpolate_3d_rect",
+    "PreComputedInterpolate3DRect",
+    "precomputed_interpolate_3d_rect",
+    "PreComputedInterpolate3DScattered",
+    "precomputed_interpolate_3d_scattered",
+    "EfieldInterpolate3DScattered",
+    "efield_interpolate_3d_scattered",
+    "euler",
+    "eulerv1",
+    "rk1",
+    "rk2",
+    "rk4",
+    "dufort_frankel",
+    "dufort_frankel_homogeneous",
+    "bwd_euler_sc",
+    "bwd_euler_ub",
+    "bwd_euler_bt",
+    "dhs",
+    "df",
+    "dfh",
+    "scnv",
+]
+
+
+def _cache_cpu_isa_list():
+    """
+    Checks for a cached CPU ISA string. If not found, runs the slow
+    detection and caches the result for future runs.
+    """
+    # Use a standard cache location.
+    cache_dir = Path.home() / ".cache" / "dendra"
+    cache_file = cache_dir / f"cpu_isa_list_{torch.__version__}"
+
+    def get_valid_vec_isa_list():
+        with open(cache_file, "rb") as f:
+            return pickle.load(f)
+
+    # --- Cache Hit ---
+    if cache_file.exists():
+        # Check if the PyTorch version matches the one used to create the cache.
+        torch._inductor.cpu_vec_isa.valid_vec_isa_list = get_valid_vec_isa_list
+        return
+
+    # --- Cache Miss or Stale Cache ---
+    print("Dendra: Performing one-time CPU capability check. This may take a minute...")
+    try:
+        os.makedirs(cache_dir, exist_ok=True)
+        start = time.time()
+        valid_vec_isa_list = torch._inductor.cpu_vec_isa.valid_vec_isa_list()
+        print(f"Dendra: CPU capability check took {(time.time() - start):.3f}s.")
+        with open(cache_file, "wb") as f:
+            pickle.dump(valid_vec_isa_list, f)
+    except Exception as e:
+        warnings.warn(
+            f"Dendra: CPU capability check failed: {e}. Falling back to default behavior."
+        )
+        return
+
+
+if _torchinductor_cache_opt_in and _should_cache_cpu_isa_for_dendra():
+    _cache_cpu_isa_list()
+
+# setup environment
+allow_tf32(bool(TF32))

dendra/__main__.py

@@ -0,0 +1,5 @@
+"""Run Dendra's command-line interface with ``python -m dendra``."""
+
+from .cli import main
+
+raise SystemExit(main())

dendra/_bootstrap.py

@@ -0,0 +1,217 @@
+# dendra/_bootstrap.py
+
+from __future__ import annotations
+
+import atexit
+import getpass
+import os
+import shutil
+import sys
+import tempfile
+import uuid
+import warnings
+from contextlib import contextmanager
+from pathlib import Path
+
+_DISABLED_INDUCTOR_CACHE_POLICIES = {
+    "0",
+    "false",
+    "none",
+    "off",
+    "disable",
+    "disabled",
+}
+_PROCESS_INDUCTOR_CACHE_POLICIES = {
+    "process",
+    "isolated",
+    "per_process",
+    "per-process",
+}
+_SHARED_INDUCTOR_CACHE_POLICIES = {"shared", "global"}
+
+
+@contextmanager
+def torch_compiler_warning_context():
+    """Filter only known PyTorch compiler deprecations in strict-warning runs."""
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore",
+            message=r".*torch\.jit\.script_method.*deprecated.*",
+            category=DeprecationWarning,
+            module=r"torch\.jit\._script",
+        )
+        warnings.filterwarnings(
+            "ignore",
+            message=r".*torch\.jit\.script.*deprecated.*",
+            category=DeprecationWarning,
+            module=r"torch\.jit\._script",
+        )
+        warnings.filterwarnings(
+            "ignore",
+            message=(
+                r".*The \.grad attribute of a Tensor that is not a leaf Tensor "
+                r"is being accessed.*"
+            ),
+            category=UserWarning,
+        )
+        warnings.filterwarnings(
+            "ignore",
+            message=(
+                r".*torch\.autograd\.function\.Function.*should not be "
+                r"instantiated\..*Methods on autograd functions.*"
+            ),
+            category=DeprecationWarning,
+        )
+        yield
+
+
+def reset_torch_compiler(*, prefer_public: bool = True) -> None:
+    """Reset Torch compiler state without leaking PyTorch-internal warnings.
+
+    PyTorch 2.13 can emit deprecations for ``torch.jit.script_method`` while
+    lazily initializing Inductor and for internal ``autograd.Function``
+    handling while tracing custom functions. Dendra still treats every other
+    warning as actionable; only those upstream deprecations are filtered around
+    compiler operations that trigger them.
+    """
+    import torch
+
+    with torch_compiler_warning_context():
+        if not prefer_public and hasattr(torch, "_dynamo"):
+            torch._dynamo.reset()
+        elif hasattr(torch, "compiler") and hasattr(torch.compiler, "reset"):
+            torch.compiler.reset()
+        elif hasattr(torch, "_dynamo") and hasattr(torch._dynamo, "reset"):
+            torch._dynamo.reset()
+
+
+def _truthy(x: str | None) -> bool:
+    return str(x).lower() in {"1", "true", "yes", "on"}
+
+
+def _should_configure_torchinductor_cache_on_import() -> bool:
+    """Return whether importing Dendra should configure TorchInductor's cache."""
+    return _truthy(os.environ.get("DENDRA_CONFIGURE_TORCHINDUCTOR_CACHE", "0"))
+
+
+def _should_cache_cpu_isa_for_dendra() -> bool:
+    """Return whether the selected policy enables Dendra's CPU ISA cache."""
+    policy = os.environ.get("DENDRA_INDUCTOR_CACHE_POLICY", "process").lower()
+    return policy in (
+        _PROCESS_INDUCTOR_CACHE_POLICIES | _SHARED_INDUCTOR_CACHE_POLICIES
+    )
+
+
+def _get_notebook_or_process_id() -> str:
+    """
+    Stable-enough process identifier for interactive notebooks and scripts.
+
+    We include PID to avoid cross-process cache sharing.
+    UUID avoids stale collisions if a PID is reused.
+    """
+    user = "unknown"
+    try:
+        user = getpass.getuser()
+    except Exception:
+        pass
+
+    return f"{user}_{os.getpid()}_{uuid.uuid4().hex[:8]}"
+
+
+def configure_torchinductor_cache_for_dendra() -> None:
+    """
+    Configure TorchInductor cache behavior before importing torch.
+
+    Environment variables
+    ---------------------
+    DENDRA_CONFIGURE_TORCHINDUCTOR_CACHE:
+        "1" opts ``import dendra`` into Dendra-managed TorchInductor cache
+        behavior, including this function and the CPU ISA cache. Default: "0".
+        The flag does not affect explicit calls to this function.
+
+    DENDRA_INDUCTOR_CACHE_POLICY:
+        "process"  -> per-process cache directory, safest for notebooks
+        "shared"   -> leave TORCHINDUCTOR_CACHE_DIR alone unless user set it
+        "off"      -> do nothing
+
+    DENDRA_INDUCTOR_CACHE_ROOT:
+        Parent directory for Dendra-owned process caches.
+        Default: tempfile.gettempdir()/dendra_torchinductor
+
+    DENDRA_INDUCTOR_DISABLE_PCH:
+        "1" disables C++ precompiled header caching where supported.
+
+    DENDRA_INDUCTOR_COMPILE_THREADS:
+        Sets TORCHINDUCTOR_COMPILE_THREADS. Default "1" in process mode.
+
+    DENDRA_INDUCTOR_CLEANUP:
+        "1" removes the Dendra-owned per-process cache at interpreter exit.
+    """
+    policy = os.environ.get("DENDRA_INDUCTOR_CACHE_POLICY", "process").lower()
+
+    if policy in _DISABLED_INDUCTOR_CACHE_POLICIES:
+        return
+
+    # This is the important limitation. If torch is already imported, some
+    # Inductor config/cache state may already be initialized.
+    if "torch" in sys.modules:
+        warnings.warn(
+            "Dendra: torch was imported before Dendra configured TorchInductor. "
+            "Dendra cannot reliably isolate TORCHINDUCTOR_CACHE_DIR after torch import. "
+            "For concurrent notebook use, import dendra before torch or set "
+            "TORCHINDUCTOR_CACHE_DIR before starting Python.",
+            RuntimeWarning,
+            stacklevel=2,
+        )
+        return
+
+    if policy in _PROCESS_INDUCTOR_CACHE_POLICIES:
+        root = Path(
+            os.environ.get(
+                "DENDRA_INDUCTOR_CACHE_ROOT",
+                str(Path(tempfile.gettempdir()) / "dendra_torchinductor"),
+            )
+        ).expanduser()
+
+        session_id = os.environ.get("DENDRA_INDUCTOR_SESSION_ID")
+        if not session_id:
+            session_id = _get_notebook_or_process_id()
+
+        cache_dir = root / session_id
+        cache_dir.mkdir(parents=True, exist_ok=True)
+
+        # Only set if user did not explicitly choose a cache directory.
+        os.environ.setdefault("TORCHINDUCTOR_CACHE_DIR", str(cache_dir))
+
+        # This is the key workaround for the malformed .h/.h.pch failure mode.
+        # It is an internal/implementation-dependent knob, so setdefault keeps
+        # user choice authoritative.
+        if _truthy(os.environ.get("DENDRA_INDUCTOR_DISABLE_PCH", "1")):
+            os.environ.setdefault("TORCHINDUCTOR_CPP_CACHE_PRECOMPILE_HEADERS", "0")
+
+        # Reduces within-process async compile races and makes failures easier
+        # to reproduce/debug. Users can override.
+        os.environ.setdefault(
+            "TORCHINDUCTOR_COMPILE_THREADS",
+            os.environ.get("DENDRA_INDUCTOR_COMPILE_THREADS", "1"),
+        )
+
+        if _truthy(os.environ.get("DENDRA_INDUCTOR_CLEANUP", "0")):
+
+            def _cleanup_cache_dir(path: Path = cache_dir) -> None:
+                shutil.rmtree(path, ignore_errors=True)
+
+            atexit.register(_cleanup_cache_dir)
+
+    elif policy in _SHARED_INDUCTOR_CACHE_POLICIES:
+        # Respect user's chosen shared cache. Do not clean it automatically.
+        if _truthy(os.environ.get("DENDRA_INDUCTOR_DISABLE_PCH", "1")):
+            os.environ.setdefault("TORCHINDUCTOR_CPP_CACHE_PRECOMPILE_HEADERS", "0")
+
+    else:
+        warnings.warn(
+            f"Dendra: unknown DENDRA_INDUCTOR_CACHE_POLICY={policy!r}; "
+            "expected 'process', 'shared', or 'off'.",
+            RuntimeWarning,
+            stacklevel=2,
+        )

dendra/cli.py

@@ -0,0 +1,51 @@
+"""Command-line entry points for Dendra deployment diagnostics."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from typing import Optional, Sequence
+
+from .diagnostics import collect_doctor_report
+
+
+def _parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog="dendra")
+    subparsers = parser.add_subparsers(dest="command")
+    doctor = subparsers.add_parser(
+        "doctor", help="inspect CPU, CUDA, and optional native-extension readiness"
+    )
+    doctor.add_argument(
+        "--require-cuda",
+        action="store_true",
+        help="treat unavailable CUDA/PyTorch CUDA support as an error",
+    )
+    doctor.add_argument(
+        "--probe-native-bitpack",
+        action="store_true",
+        help="opt into JIT build/load and tiny native bitpack kernel launches",
+    )
+    doctor.add_argument(
+        "--json",
+        action="store_true",
+        help="emit the structured report as JSON",
+    )
+    return parser
+
+
+def main(argv: Optional[Sequence[str]] = None) -> int:
+    parser = _parser()
+    args = parser.parse_args(argv)
+    if args.command != "doctor":
+        parser.print_help()
+        return 2
+
+    report = collect_doctor_report(
+        require_cuda=args.require_cuda,
+        probe_native_bitpack=args.probe_native_bitpack,
+    )
+    if args.json:
+        print(json.dumps(report.as_dict(), indent=2, sort_keys=True))
+    else:
+        print(report.to_text())
+    return 0 if report.ok else 1

dendra/const.py

@@ -0,0 +1,34 @@
+"""dendra.const
+
+Mathematical and physical constants used throughout Dendra.
+
+All values are plain Python ``float`` constants (TorchScript-friendly).
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+
+__all__ = ["PI", "E", "TAU", "R", "FARADAY"]
+
+PI: torch.jit.Final[float] = math.pi
+""":math:`\\pi`, ratio of a circle's circumference to its diameter."""
+
+E: torch.jit.Final[float] = math.e
+"""Euler's number."""
+
+TAU: torch.jit.Final[float] = math.tau
+""":math:`\\tau = 2\\pi`."""
+
+R: torch.jit.Final[float] = 8314.46261815324
+"""Molar gas constant in Dendra's millivolt convention.
+
+The value has units mV·C/(mol·K), so ``R * temperature / FARADAY`` is in mV
+when :data:`FARADAY` is expressed in C/mol. It is numerically equal to the gas
+constant in J/(kmol·K).
+"""
+
+FARADAY: torch.jit.Final[float] = 96485.33212331001
+"""Faraday constant in SI units of C/mol."""

dendra/data/__init__.py

@@ -0,0 +1,5 @@
+from .h5 import H5Reader
+
+__all__ = [
+    "H5Reader",
+]