xpcsjax 0.1.0__py3-none-any.whl

xpcsjax/__init__.py

@@ -0,0 +1,166 @@
+"""xpcsjax — unified JAX-native XPCS NLSQ fitting.
+
+Public API (lazy-loaded — heavy deps like JAX import on first use):
+
+    from xpcsjax import load_xpcs_data, fit_nlsq, ConfigManager
+
+    data = load_xpcs_data("config.yaml")
+    result = fit_nlsq(data, "config.yaml")
+    print(result.parameters)
+    result.save("output/")
+
+Env setup at import time is mirrored verbatim from homodyne/__init__.py.
+"""
+
+from __future__ import annotations
+
+# ============================================================================
+# Standard library imports
+# ============================================================================
+import importlib
+import logging
+import os
+
+# ============================================================================
+# JAX CPU Device Configuration (MUST be set before JAX import)
+# ============================================================================
+# Mirrored from homodyne/__init__.py, with one xpcsjax-specific adaptation
+# (concurrency gating, see _xla_host_device_count):
+#   - xla_force_host_platform_device_count: enables parallel evaluation paths
+#   - xla_disable_hlo_passes=constant_folding: prevents > 1 s slow-compilation
+#     warnings on HYBRID_STREAMING strategy (23M+ points) where data arrays
+#     are captured in JIT closures. Performance impact: minimal (< 5 ms/call).
+
+
+def _detect_worker_count() -> int:
+    """Concurrent fit-process count from the pool / pytest-xdist env-vars (>=1).
+
+    Inlined here (not imported from ``optimization.nlsq.memory``) on purpose:
+    this runs *before* the first JAX import, and importing that module could pull
+    JAX in early and defeat the env-before-import ordering this block exists for.
+    """
+    for _env_var in ("XPCSJAX_FIT_CONCURRENCY", "PYTEST_XDIST_WORKER_COUNT"):
+        _raw = os.environ.get(_env_var)
+        if _raw:
+            try:
+                return max(1, int(_raw))
+            except ValueError:
+                pass
+    return 1
+
+
+def _xla_host_device_count(worker_count: int) -> int:
+    """Return the number of host CPU devices to force for XLA.
+
+    A lone fit benefits from 4 host devices (parallel evaluation paths). But
+    under parallelism (pytest-xdist or the production multistart / accumulator
+    pools) each of N worker processes forcing 4 devices means 4*N logical devices
+    contending for the physical cores — wasted per-worker compile/buffer overhead
+    and RAM. So drop to a single device per process whenever more than one fit
+    runs concurrently.
+    """
+    return 4 if worker_count <= 1 else 1
+
+
+_WORKER_COUNT = _detect_worker_count()
+_DEFAULT_XLA_FLAGS = [
+    f"--xla_force_host_platform_device_count={_xla_host_device_count(_WORKER_COUNT)}",
+    "--xla_disable_hlo_passes=constant_folding",
+]
+
+# JAX must be in float64 for parameters spanning 6+ orders of magnitude.
+# This env var MUST be set BEFORE the first JAX import.
+os.environ.setdefault("JAX_ENABLE_X64", "1")
+
+if "XLA_FLAGS" not in os.environ:
+    os.environ["XLA_FLAGS"] = " ".join(_DEFAULT_XLA_FLAGS)
+else:
+    existing = os.environ["XLA_FLAGS"]
+    flags_to_add = []
+    for flag in _DEFAULT_XLA_FLAGS:
+        flag_name = flag.split("=")[0]
+        if flag_name not in existing:
+            flags_to_add.append(flag)
+    if flags_to_add:
+        os.environ["XLA_FLAGS"] += " " + " ".join(flags_to_add)
+
+# Pin JAX to CPU (CPU-only; no GPU support). Setting this at
+# package import time (before any jax import) is the *only* place this works
+# reliably — spawn-pool worker init runs *after* the worker's `import jax`,
+# which is why xpcsjax.viz.nlsq_plots.{_worker_init_cpu_only,_render_one_angle_worker}
+# can't set this themselves. Child processes inherit os.environ from the parent.
+os.environ.setdefault("JAX_PLATFORMS", "cpu")
+os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
+os.environ.setdefault("XLA_PYTHON_CLIENT_ALLOCATOR", "platform")
+
+# Cap each worker's JAX arena under parallelism so one fit process can't claim
+# the whole device. No-op on the CPU backend (the v0.1 target), but a correct
+# per-worker bound for any future GPU backend; setdefault leaves an explicit
+# user/operator override untouched.
+if _WORKER_COUNT > 1:
+    os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", f"{max(0.05, 0.9 / _WORKER_COUNT):.4f}")
+
+# Suppress NLSQ GPU warnings (CPU-only; no GPU support)
+os.environ.setdefault("NLSQ_SKIP_GPU_CHECK", "1")
+
+# Suppress JAX backend logs (GPU fallback warnings on CPU-only systems)
+logging.getLogger("jax._src.xla_bridge").setLevel(logging.ERROR)
+logging.getLogger("jax._src.compiler").setLevel(logging.ERROR)
+
+# ============================================================================
+# Version
+# ============================================================================
+__version__ = "0.1.0"
+
+# ============================================================================
+# Lazy public API
+# ============================================================================
+_LAZY_EXPORTS = {
+    "load_xpcs_data": "xpcsjax.data",
+    "fit_nlsq": "xpcsjax.optimization.nlsq",
+    "ConfigManager": "xpcsjax.config",
+    "generate_nlsq_plots": "xpcsjax.viz",
+    "HomodyneModel": "xpcsjax.core",
+    "HeterodyneModel": "xpcsjax.core",
+    "OptimizationResult": "xpcsjax.optimization.nlsq.results",
+}
+
+# TYPE_CHECKING block for IDE / Pyright static visibility. All submodules
+# below now export their public symbol, so the original deferral comment
+# (Tasks 6/11/15/19/20/28) is resolved.
+from typing import TYPE_CHECKING as _TYPE_CHECKING
+
+if _TYPE_CHECKING:
+    from xpcsjax.config import ConfigManager
+    from xpcsjax.core import HeterodyneModel, HomodyneModel
+    from xpcsjax.data import load_xpcs_data
+    from xpcsjax.optimization.nlsq import fit_nlsq
+    from xpcsjax.optimization.nlsq.results import OptimizationResult
+    from xpcsjax.viz import generate_nlsq_plots
+
+
+def __getattr__(name: str):  # noqa: D401
+    """Lazy attribute loader for the documented public API."""
+    if name in _LAZY_EXPORTS:
+        module = importlib.import_module(_LAZY_EXPORTS[name])
+        attr = getattr(module, name)
+        globals()[name] = attr
+        return attr
+    raise AttributeError(f"module 'xpcsjax' has no attribute {name!r}")
+
+
+# Literal __all__ for Pyright's reportUnsupportedDunderAll; kept in sync
+# with _LAZY_EXPORTS by the runtime assertion below.
+__all__ = [
+    "load_xpcs_data",
+    "fit_nlsq",
+    "ConfigManager",
+    "generate_nlsq_plots",
+    "HomodyneModel",
+    "HeterodyneModel",
+    "OptimizationResult",
+]
+
+assert set(__all__) == set(_LAZY_EXPORTS), (
+    "xpcsjax public API mismatch between __all__ and _LAZY_EXPORTS"
+)

xpcsjax/cli/__init__.py

@@ -0,0 +1,78 @@
+"""Command-line interface for xpcsjax (NLSQ-only XPCS analysis).
+
+This subpackage is lazy-loaded: importing :mod:`xpcsjax.cli` does NOT
+import JAX or any of the heavy submodules. Attribute access (e.g.
+``xpcsjax.cli.main``) triggers the actual import via ``__getattr__``.
+
+Mirrors heterodyne's CLI lazy-import surface so test mocks can target
+the same import paths.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+# TYPE_CHECKING block: gives Pyright / mypy / IDEs static visibility of the
+# lazy-exported symbols without paying the import cost at runtime. The
+# ``__getattr__`` below is what actually resolves them at first access.
+if TYPE_CHECKING:
+    from xpcsjax.cli.args_parser import create_parser, validate_args
+    from xpcsjax.cli.commands import dispatch_command
+    from xpcsjax.cli.config_generator import main as config_main
+    from xpcsjax.cli.config_handling import apply_cli_overrides, load_and_merge_config
+    from xpcsjax.cli.data_pipeline import load_and_validate_data, resolve_phi_angles
+    from xpcsjax.cli.main import main
+    from xpcsjax.cli.optimization_runner import run_nlsq
+    from xpcsjax.cli.plot_dispatch import dispatch_plots
+    from xpcsjax.cli.xla_config import configure_xla
+
+
+_IMPORTS: dict[str, tuple[str, str]] = {
+    "main": ("xpcsjax.cli.main", "main"),
+    "config_main": ("xpcsjax.cli.config_generator", "main"),
+    "configure_xla": ("xpcsjax.cli.xla_config", "configure_xla"),
+    "create_parser": ("xpcsjax.cli.args_parser", "create_parser"),
+    "validate_args": ("xpcsjax.cli.args_parser", "validate_args"),
+    "dispatch_command": ("xpcsjax.cli.commands", "dispatch_command"),
+    "load_and_merge_config": (
+        "xpcsjax.cli.config_handling",
+        "load_and_merge_config",
+    ),
+    "apply_cli_overrides": (
+        "xpcsjax.cli.config_handling",
+        "apply_cli_overrides",
+    ),
+    "load_and_validate_data": (
+        "xpcsjax.cli.data_pipeline",
+        "load_and_validate_data",
+    ),
+    "resolve_phi_angles": ("xpcsjax.cli.data_pipeline", "resolve_phi_angles"),
+    "run_nlsq": ("xpcsjax.cli.optimization_runner", "run_nlsq"),
+    "dispatch_plots": ("xpcsjax.cli.plot_dispatch", "dispatch_plots"),
+}
+
+
+def __getattr__(name: str) -> Any:
+    if name in _IMPORTS:
+        module_path, attr = _IMPORTS[name]
+        import importlib
+
+        module = importlib.import_module(module_path)
+        return getattr(module, attr)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    "main",
+    "config_main",
+    "configure_xla",
+    "create_parser",
+    "validate_args",
+    "dispatch_command",
+    "load_and_merge_config",
+    "apply_cli_overrides",
+    "load_and_validate_data",
+    "resolve_phi_angles",
+    "run_nlsq",
+    "dispatch_plots",
+]

xpcsjax/cli/args_parser.py

@@ -0,0 +1,442 @@
+"""Argument parser for the xpcsjax CLI.
+
+NLSQ-only by design (see project CLAUDE.md); Bayesian sampling flags from
+the upstream heterodyne CLI are intentionally absent.
+
+Parameter override flags map to canonical names per
+``parameter_registry._MODE_PARAMS``:
+
+* ``static_anisotropic`` / ``static_isotropic`` — D0, alpha, D_offset
+* ``laminar_flow`` — D0, alpha, D_offset, gamma_dot_t0, beta,
+  gamma_dot_t_offset, phi0
+* ``two_component`` — D0_ref/alpha_ref/D_offset_ref,
+  D0_sample/alpha_sample/D_offset_sample, v0, v_beta, v_offset,
+  f0..f3, phi0_het. The reference/sample transport params and phi0_het
+  have no CLI flag (14 params is too many for the command line); set
+  them in the YAML config. The v0/v_beta/v_offset/f0..f3 flags below DO
+  apply to two_component.
+
+Flags whose canonical name is not in the active mode's parameter set are
+silently ignored by ``config_handling.apply_cli_overrides`` (it
+intersects against ``ConfigManager.get_active_parameters()`` before
+writing the canonical ``initial_parameters`` block).
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+_VALID_MODES = ("static_anisotropic", "static_isotropic", "laminar_flow", "two_component")
+
+
+def create_parser() -> argparse.ArgumentParser:
+    """Build the xpcsjax CLI argument parser."""
+    parser = argparse.ArgumentParser(
+        prog="xpcsjax",
+        description=(
+            "xpcsjax — JAX-native NLSQ fitting for homodyne / heterodyne XPCS. "
+            "Bayesian sampling is permanently out of scope for this package; "
+            "use the upstream homodyne / heterodyne packages for that."
+        ),
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Run NLSQ fit with a YAML config
+  xpcsjax --config analysis.yaml
+
+  # Override the output directory
+  xpcsjax --config analysis.yaml --output ./results
+
+  # Multistart NLSQ with 16 restarts
+  xpcsjax --config analysis.yaml --multistart --multistart-n 16
+
+  # Plot experimental data only (skip optimization)
+  xpcsjax --config analysis.yaml --plot-experimental-data
+
+  # Plot simulated C2 heatmaps from config parameters
+  xpcsjax --config analysis.yaml --plot-simulated-data --phi-angles 0,45,90,135
+
+  # Launch the interactive analysis workbench (GUI) — separate console script
+  xpcsjax-gui
+
+Exit codes:
+  0   Analysis completed and the optimizer converged (or no convergence
+      check applies, e.g. plot-only runs).
+  1   Unhandled exception during analysis. See log for traceback.
+  2   Analysis ran but the optimizer did NOT converge. Output files are
+      still written but the fit is not trustworthy.
+  130 Interrupted by the user (Ctrl-C).
+""",
+    )
+
+    # ------------------------------------------------------------------
+    # Required: config path
+    # ------------------------------------------------------------------
+    parser.add_argument(
+        "--config",
+        "-c",
+        type=Path,
+        required=True,
+        help="Path to YAML configuration file (required).",
+    )
+
+    # ------------------------------------------------------------------
+    # Output
+    # ------------------------------------------------------------------
+    parser.add_argument(
+        "--output",
+        "-o",
+        type=Path,
+        default=None,
+        help="Output directory (overrides ``output.directory`` in YAML).",
+    )
+    parser.add_argument(
+        "--output-format",
+        choices=["json", "npz", "both"],
+        default="both",
+        help="Format for saved results (default: both).",
+    )
+
+    # ------------------------------------------------------------------
+    # Mode / phi
+    # ------------------------------------------------------------------
+    parser.add_argument(
+        "--mode",
+        choices=_VALID_MODES,
+        default=None,
+        help=(f"Force ``analysis_mode`` (overrides YAML). Must be one of {_VALID_MODES}."),
+    )
+    parser.add_argument(
+        "--phi",
+        type=float,
+        nargs="+",
+        default=None,
+        help="Phi angles to analyze, in degrees (overrides config).",
+    )
+
+    # ------------------------------------------------------------------
+    # NLSQ options
+    # ------------------------------------------------------------------
+    nlsq_group = parser.add_argument_group("NLSQ options", "Solver and multistart controls.")
+    # store_const with default=None (not store_true) so the override layer
+    # can distinguish "user did not pass --multistart" (None → leave YAML
+    # untouched) from "user passed it" (True). store_true's False default
+    # would clobber a YAML ``multi_start.enable: true`` on every run.
+    nlsq_group.add_argument(
+        "--multistart",
+        action="store_const",
+        const=True,
+        default=None,
+        help="Enable LHS multistart for NLSQ (overrides YAML).",
+    )
+    nlsq_group.add_argument(
+        "--no-multistart",
+        dest="multistart",
+        action="store_const",
+        const=False,
+        help="Explicitly disable multistart (overrides a YAML enable).",
+    )
+    nlsq_group.add_argument(
+        "--multistart-n",
+        type=int,
+        default=None,
+        help="Number of multistart restarts (default: from config, else 10).",
+    )
+    nlsq_group.add_argument(
+        "--max-iterations",
+        type=int,
+        default=None,
+        help="Maximum trust-region iterations (overrides config).",
+    )
+    nlsq_group.add_argument(
+        "--tolerance",
+        type=float,
+        default=None,
+        help="NLSQ convergence tolerance (overrides config).",
+    )
+
+    # ------------------------------------------------------------------
+    # Parameter overrides (highest precedence)
+    # CLI > YAML > parameter_registry defaults.
+    # Covers the union of all four modes; mode-incompatible flags are
+    # silently ignored downstream.
+    # ------------------------------------------------------------------
+    param_group = parser.add_argument_group(
+        "parameter overrides",
+        "Override initial parameter values (highest precedence). "
+        "Flags for parameters not used by the active mode are ignored.",
+    )
+    # Core transport (all modes)
+    param_group.add_argument(
+        "--initial-D0",
+        type=float,
+        default=None,
+        metavar="VAL",
+        help="Diffusion prefactor D0 [Ų/s^α].",
+    )
+    param_group.add_argument(
+        "--initial-alpha",
+        type=float,
+        default=None,
+        metavar="VAL",
+        help="Transport exponent alpha.",
+    )
+    param_group.add_argument(
+        "--initial-D-offset",
+        type=float,
+        default=None,
+        metavar="VAL",
+        help="Transport offset D_offset [Ų/s].",
+    )
+    # Laminar flow / two-component velocity
+    param_group.add_argument(
+        "--initial-gamma-dot-t0",
+        type=float,
+        default=None,
+        metavar="VAL",
+        help="Shear rate prefactor (laminar_flow).",
+    )
+    param_group.add_argument(
+        "--initial-gamma-dot-t-offset",
+        type=float,
+        default=None,
+        metavar="VAL",
+        help="Shear rate offset (laminar_flow).",
+    )
+    param_group.add_argument(
+        "--initial-beta",
+        type=float,
+        default=None,
+        metavar="VAL",
+        help="Velocity exponent beta (laminar_flow).",
+    )
+    param_group.add_argument(
+        "--initial-v-beta",
+        type=float,
+        default=None,
+        metavar="VAL",
+        help="Velocity exponent v_beta (two_component).",
+    )
+    param_group.add_argument(
+        "--initial-v0",
+        type=float,
+        default=None,
+        metavar="VAL",
+        help="Velocity prefactor v0 (two_component).",
+    )
+    param_group.add_argument(
+        "--initial-v-offset",
+        type=float,
+        default=None,
+        metavar="VAL",
+        help="Velocity offset v_offset (two_component).",
+    )
+    # Angle parameters
+    param_group.add_argument(
+        "--initial-phi0",
+        type=float,
+        default=None,
+        metavar="VAL",
+        help="Flow angle offset phi0 [degrees].",
+    )
+    # Two-component Fourier amplitudes (per-angle fraction)
+    param_group.add_argument(
+        "--initial-f0",
+        type=float,
+        default=None,
+        metavar="VAL",
+        help="Sample fraction amplitude f0 (two_component).",
+    )
+    param_group.add_argument(
+        "--initial-f1",
+        type=float,
+        default=None,
+        metavar="VAL",
+        help="Fourier coefficient f1 (two_component).",
+    )
+    param_group.add_argument(
+        "--initial-f2",
+        type=float,
+        default=None,
+        metavar="VAL",
+        help="Fourier coefficient f2 (two_component).",
+    )
+    param_group.add_argument(
+        "--initial-f3",
+        type=float,
+        default=None,
+        metavar="VAL",
+        help="Fourier coefficient f3 (two_component).",
+    )
+    # Note: per-angle scaling (contrast/offset) is not a single-value CLI
+    # override — there is one pair per phi angle. Set scaling in the YAML
+    # config's per-angle scaling block instead.
+
+    # ------------------------------------------------------------------
+    # Verbosity
+    # ------------------------------------------------------------------
+    parser.add_argument(
+        "--verbose",
+        "-v",
+        action="count",
+        default=0,
+        help="Increase verbosity (-v, -vv).",
+    )
+    parser.add_argument(
+        "--quiet",
+        "-q",
+        action="store_true",
+        help="Suppress all output except errors.",
+    )
+
+    # ------------------------------------------------------------------
+    # Performance
+    # ------------------------------------------------------------------
+    parser.add_argument(
+        "--threads",
+        type=int,
+        default=None,
+        help="CPU thread count for XLA (default: auto).",
+    )
+    parser.add_argument(
+        "--no-jit",
+        action="store_true",
+        help="Disable JIT compilation (for debugging only — much slower).",
+    )
+
+    # ------------------------------------------------------------------
+    # Plotting
+    # ------------------------------------------------------------------
+    plot_group = parser.add_mutually_exclusive_group()
+    plot_group.add_argument(
+        "--plot",
+        dest="plot",
+        action="store_true",
+        default=True,
+        help="Generate plots after fitting (default).",
+    )
+    plot_group.add_argument(
+        "--no-plot",
+        dest="plot",
+        action="store_false",
+        help="Skip plot generation.",
+    )
+    parser.add_argument(
+        "--save-plots",
+        action="store_true",
+        help="Save fit-comparison plots to the output directory.",
+    )
+    parser.add_argument(
+        "--plotting-backend",
+        choices=["auto", "matplotlib", "datashader"],
+        default="auto",
+        help=(
+            "Plotting backend: auto (Datashader if installed), "
+            "matplotlib, or datashader (default: %(default)s)."
+        ),
+    )
+    parser.add_argument(
+        "--parallel-plots",
+        action="store_true",
+        help="Generate plots in parallel via multiprocessing (Datashader path).",
+    )
+    parser.add_argument(
+        "--phi-angles",
+        type=str,
+        default=None,
+        help=("Comma-separated phi angles in degrees for simulated data (e.g. '0,45,90,135')."),
+    )
+
+    # Standalone plot modes (skip optimization)
+    parser.add_argument(
+        "--plot-experimental-data",
+        action="store_true",
+        help="Plot experimental data for QC (skip optimization).",
+    )
+    parser.add_argument(
+        "--plot-simulated-data",
+        action="store_true",
+        help="Plot simulated C2 heatmaps from config parameters (skip optimization).",
+    )
+    parser.add_argument(
+        "--contrast",
+        type=float,
+        default=0.3,
+        help="Contrast for simulated data (default: %(default)s; requires --plot-simulated-data).",
+    )
+    parser.add_argument(
+        "--offset-sim",
+        type=float,
+        default=1.0,
+        help="Offset for simulated data (default: %(default)s; requires --plot-simulated-data).",
+    )
+
+    # ------------------------------------------------------------------
+    # Version
+    # ------------------------------------------------------------------
+    _add_version_arg(parser)
+
+    return parser
+
+
+def _add_version_arg(parser: argparse.ArgumentParser) -> None:
+    """Add ``--version`` with a best-effort version resolution."""
+    try:
+        import importlib.metadata as _md
+
+        version = _md.version("xpcsjax")
+    except Exception:  # pragma: no cover — uninstalled / dev tree
+        try:
+            from xpcsjax import __version__ as version
+        except Exception:
+            version = "unknown"
+    parser.add_argument(
+        "--version",
+        action="version",
+        version=f"%(prog)s {version}",
+    )
+
+
+def validate_args(args: argparse.Namespace) -> list[str]:
+    """Light validation pass. Returns non-fatal warning strings.
+
+    Raises ``FileNotFoundError`` for unrecoverable issues (missing config).
+    """
+    warnings: list[str] = []
+
+    if not args.config.exists():
+        raise FileNotFoundError(f"Configuration file not found: {args.config}")
+
+    if args.verbose > 0 and args.quiet:
+        warnings.append("Both --verbose and --quiet specified; --quiet wins.")
+        args.verbose = 0
+
+    phi_angles_str: str | None = getattr(args, "phi_angles", None)
+    if phi_angles_str is not None:
+        try:
+            [float(x.strip()) for x in phi_angles_str.split(",")]
+        except ValueError:
+            warnings.append(
+                f"--phi-angles must be comma-separated numbers "
+                f"(e.g. '0,45,90,135'); got: {phi_angles_str!r}"
+            )
+
+    if args.plot_experimental_data and args.plot_simulated_data:
+        warnings.append(
+            "Both --plot-experimental-data and --plot-simulated-data given; running both passes."
+        )
+
+    if args.multistart_n is not None and args.multistart_n <= 0:
+        warnings.append(
+            f"--multistart-n must be positive; got {args.multistart_n}. "
+            "Falling back to config value."
+        )
+        args.multistart_n = None
+
+    return warnings
+
+
+build_parser = create_parser  # uniform factory name used by completion_spec
+
+__all__ = ["build_parser", "create_parser", "validate_args"]