openscvx 0.3.2.dev170__py3-none-any.whl

openscvx/utils/cache.py

@@ -0,0 +1,131 @@
+"""Cache management for compiled solvers.
+
+This module provides utilities for managing the cache directory where compiled
+JAX solvers are stored. The cache location follows platform conventions:
+
+- **Linux**: ``~/.cache/openscvx/``
+- **macOS**: ``~/Library/Caches/openscvx/``
+- **Windows**: ``%LOCALAPPDATA%/openscvx/Cache/``
+
+The cache location can be overridden by setting the ``OPENSCVX_CACHE_DIR``
+environment variable.
+
+Example:
+    Get the cache directory::
+
+        import openscvx as ox
+        print(ox.get_cache_dir())  # /home/user/.cache/openscvx
+
+    Clear all cached solvers::
+
+        import openscvx as ox
+        ox.clear_cache()
+
+    Check cache size::
+
+        import openscvx as ox
+        size_mb = ox.get_cache_size() / (1024 * 1024)
+        print(f"Cache size: {size_mb:.1f} MB")
+"""
+
+import os
+import shutil
+import sys
+from pathlib import Path
+
+
+def get_cache_dir() -> Path:
+    """Get the cache directory for compiled solvers.
+
+    The cache location is determined in the following order:
+    1. ``OPENSCVX_CACHE_DIR`` environment variable (if set)
+    2. Platform-specific default:
+       - Linux: ``~/.cache/openscvx/``
+       - macOS: ``~/Library/Caches/openscvx/``
+       - Windows: ``%LOCALAPPDATA%/openscvx/Cache/``
+
+    Returns:
+        Path to the cache directory (may not exist yet)
+    """
+    # Check environment variable override
+    env_dir = os.environ.get("OPENSCVX_CACHE_DIR")
+    if env_dir:
+        return Path(env_dir)
+
+    # Platform-specific defaults
+    if sys.platform == "darwin":
+        # macOS: ~/Library/Caches/openscvx/
+        return Path.home() / "Library" / "Caches" / "openscvx"
+    elif sys.platform == "win32":
+        # Windows: %LOCALAPPDATA%/openscvx/Cache/
+        local_app_data = os.environ.get("LOCALAPPDATA")
+        if local_app_data:
+            return Path(local_app_data) / "openscvx" / "Cache"
+        else:
+            # Fallback if LOCALAPPDATA not set
+            return Path.home() / "AppData" / "Local" / "openscvx" / "Cache"
+    else:
+        # Linux and others: follow XDG Base Directory Specification
+        xdg_cache = os.environ.get("XDG_CACHE_HOME")
+        if xdg_cache:
+            return Path(xdg_cache) / "openscvx"
+        else:
+            return Path.home() / ".cache" / "openscvx"
+
+
+def clear_cache() -> int:
+    """Clear all cached compiled solvers.
+
+    Removes all files in the cache directory. The directory itself is
+    preserved but emptied.
+
+    Returns:
+        Number of files deleted
+
+    Example:
+        Clear the cache::
+
+            import openscvx as ox
+            deleted = ox.clear_cache()
+            print(f"Deleted {deleted} cached files")
+    """
+    cache_dir = get_cache_dir()
+    if not cache_dir.exists():
+        return 0
+
+    count = 0
+    for item in cache_dir.iterdir():
+        if item.is_file():
+            item.unlink()
+            count += 1
+        elif item.is_dir():
+            shutil.rmtree(item)
+            count += 1
+
+    return count
+
+
+def get_cache_size() -> int:
+    """Get the total size of the cache in bytes.
+
+    Returns:
+        Total size of all files in the cache directory in bytes.
+        Returns 0 if the cache directory doesn't exist.
+
+    Example:
+        Check cache size in megabytes::
+
+            import openscvx as ox
+            size_mb = ox.get_cache_size() / (1024 * 1024)
+            print(f"Cache size: {size_mb:.1f} MB")
+    """
+    cache_dir = get_cache_dir()
+    if not cache_dir.exists():
+        return 0
+
+    total = 0
+    for item in cache_dir.rglob("*"):
+        if item.is_file():
+            total += item.stat().st_size
+
+    return total

openscvx/utils/caching.py

@@ -0,0 +1,210 @@
+import hashlib
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple
+
+import jax
+import numpy as np
+from jax import export
+
+from openscvx.utils.cache import get_cache_dir
+
+if TYPE_CHECKING:
+    from openscvx.symbolic.problem import SymbolicProblem
+
+
+def get_solver_cache_paths(
+    symbolic_problem: "SymbolicProblem",
+    dt: float,
+    total_time: float,
+    cache_dir: Optional[Path] = None,
+) -> Tuple[Path, Path]:
+    """Generate cache file paths using symbolic AST hashing.
+
+    This function computes a hash based on the symbolic structure of the problem,
+    which is more stable than hashing lowered JAX code. Two problems with the same
+    mathematical structure will produce the same hash, regardless of variable names.
+
+    Args:
+        symbolic_problem: The preprocessed SymbolicProblem
+        dt: Time step for propagation
+        total_time: Total simulation time
+        cache_dir: Directory to store cached solvers. If None, uses the default
+            cache directory (see :func:`openscvx.get_cache_dir`).
+
+    Returns:
+        Tuple of (discretization_solver_path, propagation_solver_path)
+    """
+    from openscvx.symbolic.hashing import hash_symbolic_problem
+
+    # Get the structural hash of the symbolic problem
+    problem_hash = hash_symbolic_problem(symbolic_problem)
+
+    # Include runtime config in the hash
+    final_hasher = hashlib.sha256()
+    final_hasher.update(problem_hash.encode())
+    final_hasher.update(f"dt:{dt}".encode())
+    final_hasher.update(f"total_time:{total_time}".encode())
+    final_hash = final_hasher.hexdigest()[:32]
+
+    solver_dir = cache_dir if cache_dir is not None else get_cache_dir()
+    solver_dir.mkdir(parents=True, exist_ok=True)
+
+    dis_solver_file = solver_dir / f"compiled_discretization_solver_{final_hash}.jax"
+    prop_solver_file = solver_dir / f"compiled_propagation_solver_{final_hash}.jax"
+
+    return dis_solver_file, prop_solver_file
+
+
+def load_or_compile_discretization_solver(
+    discretization_solver: callable,
+    cache_file: Path,
+    params: Dict[str, Any],
+    n_discretization_nodes: int,
+    n_states: int,
+    n_controls: int,
+    save_compiled: bool = False,
+    debug: bool = False,
+) -> callable:
+    """Load discretization solver from cache or compile and cache it.
+
+    Args:
+        discretization_solver: The solver function to compile
+        cache_file: Path to cache file
+        params: Parameters dictionary
+        n_discretization_nodes: Number of discretization nodes
+        n_states: Number of state variables
+        n_controls: Number of control variables
+        save_compiled: Whether to save/load compiled solvers
+        debug: Whether in debug mode (skip compilation)
+
+    Returns:
+        Compiled discretization solver
+    """
+    if debug:
+        return discretization_solver
+
+    if save_compiled:
+        try:
+            with open(cache_file, "rb") as f:
+                serial_dis = f.read()
+            compiled_solver = export.deserialize(serial_dis)
+            print("✓ Loaded existing discretization solver")
+            return compiled_solver
+        except FileNotFoundError:
+            print("Compiling discretization solver...")
+
+    else:
+        print("Compiling discretization solver (not saving/loading from disk)...")
+
+    # Pass parameters as a single dictionary
+    compiled_solver = export.export(jax.jit(discretization_solver))(
+        np.ones((n_discretization_nodes, n_states)),
+        np.ones((n_discretization_nodes, n_controls)),
+        params,
+    )
+
+    if save_compiled:
+        with open(cache_file, "wb") as f:
+            f.write(compiled_solver.serialize())
+        print("✓ Discretization solver compiled and saved")
+
+    return compiled_solver
+
+
+def load_or_compile_propagation_solver(
+    propagation_solver: callable,
+    cache_file: Path,
+    params: Dict[str, Any],
+    n_states_prop: int,
+    n_controls: int,
+    max_tau_len: int,
+    save_compiled: bool = False,
+) -> callable:
+    """Load propagation solver from cache or compile and cache it.
+
+    Args:
+        propagation_solver: The solver function to compile
+        cache_file: Path to cache file
+        params: Parameters dictionary
+        n_states_prop: Number of propagation state variables
+        n_controls: Number of control variables
+        max_tau_len: Maximum tau length for propagation
+        save_compiled: Whether to save/load compiled solvers
+
+    Returns:
+        Compiled propagation solver
+    """
+    if save_compiled:
+        try:
+            with open(cache_file, "rb") as f:
+                serial_prop = f.read()
+            compiled_solver = export.deserialize(serial_prop)
+            print("✓ Loaded existing propagation solver")
+            return compiled_solver
+        except FileNotFoundError:
+            print("Compiling propagation solver...")
+
+    else:
+        print("Compiling propagation solver (not saving/loading from disk)...")
+
+    # Pass parameters as a single dictionary
+    compiled_solver = export.export(jax.jit(propagation_solver))(
+        np.ones(n_states_prop),  # x_0
+        (0.0, 0.0),  # time span
+        np.ones((1, n_controls)),  # controls_current
+        np.ones((1, n_controls)),  # controls_next
+        np.ones((1, 1)),  # tau_0
+        np.ones((1, 1)).astype("int"),  # segment index
+        0,  # idx_s_stop
+        np.ones((max_tau_len,)),  # save_time (tau_cur_padded)
+        np.ones((max_tau_len,), dtype=bool),  # mask_padded (boolean mask)
+        params,  # additional parameters as dict
+    )
+
+    if save_compiled:
+        with open(cache_file, "wb") as f:
+            f.write(compiled_solver.serialize())
+        print("✓ Propagation solver compiled and saved")
+
+    return compiled_solver
+
+
+def prime_propagation_solver(
+    propagation_solver: callable, params: Dict[str, Any], settings
+) -> None:
+    """Prime the propagation solver with a test call to ensure it works.
+
+    Args:
+        propagation_solver: Compiled propagation solver
+        params: Parameters dictionary
+        settings: Settings configuration object
+    """
+    try:
+        x_0 = np.ones(settings.sim.x_prop.initial.shape, dtype=settings.sim.x_prop.initial.dtype)
+        tau_grid = (0.0, 1.0)
+        controls_current = np.ones((1, settings.sim.u.shape[0]), dtype=settings.sim.u.guess.dtype)
+        controls_next = np.ones((1, settings.sim.u.shape[0]), dtype=settings.sim.u.guess.dtype)
+        tau_init = np.array([[0.0]], dtype=np.float64)
+        node = np.array([[0]], dtype=np.int64)
+        idx_s_stop = settings.sim.time_dilation_slice.stop
+        save_time = np.ones((settings.prp.max_tau_len,), dtype=np.float64)
+        mask_padded = np.ones((settings.prp.max_tau_len,), dtype=bool)
+        # Create dummy params dict with same structure
+        dummy_params = {
+            name: np.ones_like(value) if hasattr(value, "shape") else float(value)
+            for name, value in params.items()
+        }
+        propagation_solver.call(
+            x_0,
+            tau_grid,
+            controls_current,
+            controls_next,
+            tau_init,
+            node,
+            idx_s_stop,
+            save_time,
+            mask_padded,
+            dummy_params,
+        )
+    except Exception as e:
+        print(f"[Initialization] Priming propagation_solver.call failed: {e}")

openscvx/utils/printing.py

@@ -0,0 +1,301 @@
+import queue
+import sys
+import time
+import warnings
+from importlib.metadata import PackageNotFoundError, version
+
+import jax
+import numpy as np
+from termcolor import colored
+
+from openscvx.algorithms import OptimizationResults
+
+warnings.filterwarnings("ignore")
+
+
+# Define colors for printing
+col_main = "blue"
+col_pos = "green"
+col_neg = "red"
+
+
+def get_version() -> str:
+    try:
+        return version("openscvx")
+    except PackageNotFoundError:
+        return "0.0.0"
+
+
+def print_summary_box(lines, title="Summary"):
+    """
+    Print a centered summary box with the given lines.
+
+    Args:
+        lines (list): List of strings to display in the box
+        title (str): Title for the box (default: "Summary")
+    """
+    # Find the longest line (excluding the title which will be handled separately)
+    content_lines = lines[1:] if len(lines) > 1 else []
+    max_content_width = max(len(line) for line in content_lines) if content_lines else 0
+    title_width = len(title)
+
+    # Box width should accommodate both title and content
+    box_width = max(max_content_width, title_width) + 4  # Add padding for the box borders
+
+    # Center with respect to the 89-character horizontal lines in io.py
+    total_width = 89
+    if box_width <= total_width:
+        indent = (total_width - box_width) // 2
+    else:
+        # If box is wider than 89 chars, use a smaller fixed indentation
+        indent = 2
+
+    # Print the box with dynamic width and centering
+    print(f"\n{' ' * indent}╭{'─' * box_width}╮")
+    print(f"{' ' * indent}│ {title:^{box_width - 2}} │")
+    print(f"{' ' * indent}├{'─' * box_width}┤")
+    for line in content_lines:
+        print(f"{' ' * indent}│ {line:<{box_width - 2}} │")
+    print(f"{' ' * indent}╰{'─' * box_width}╯\n")
+
+
+def print_problem_summary(settings, lowered):
+    """
+    Print the problem summary box.
+
+    Args:
+        settings: Configuration settings containing problem information
+        lowered: LoweredProblem from lower_symbolic_problem()
+    """
+    n_nodal_convex = len(lowered.cvxpy_constraints.constraints)
+    n_nodal_nonconvex = len(lowered.jax_constraints.nodal)
+    n_ctcs = len(lowered.jax_constraints.ctcs)
+    n_augmented = settings.sim.n_states - settings.sim.true_state_slice.stop
+
+    # Count CVXPy variables, parameters, and constraints
+    from openscvx.solvers import optimal_control_problem
+
+    try:
+        # Build OCP using LoweredProblem
+        prob = optimal_control_problem(settings, lowered)
+
+        # Get the actual problem size information like CVXPy verbose output
+        n_cvx_variables = sum(var.size for var in prob.variables())
+        n_cvx_parameters = sum(param.size for param in prob.parameters())
+        n_cvx_constraints = sum(constraint.size for constraint in prob.constraints)
+    except Exception:
+        # Fallback if problem construction fails
+        n_cvx_variables = 0
+        n_cvx_parameters = 0
+        n_cvx_constraints = 0
+
+    # Get JAX backend information
+    jax_backend = jax.devices()[0].platform.upper()
+    jax_version = jax.__version__
+
+    # Build weights string conditionally
+    if isinstance(settings.scp.lam_vc, np.ndarray):
+        lam_vc_str = f"λ_vc=matrix({settings.scp.lam_vc.shape})"
+    else:
+        lam_vc_str = f"λ_vc={settings.scp.lam_vc:4.1f}"
+    weights_parts = [
+        f"λ_cost={settings.scp.lam_cost:4.1f}",
+        f"λ_tr={settings.scp.w_tr:4.1f}",
+        lam_vc_str,
+    ]
+
+    # Add λ_vb only if there are nodal nonconvex constraints
+    if n_nodal_nonconvex > 0:
+        weights_parts.append(f"λ_vb={settings.scp.lam_vb:4.1f}")
+
+    weights_str = ", ".join(weights_parts)
+
+    lines = [
+        "Problem Summary",
+        (
+            f"Dimensions: {settings.sim.n_states} states ({n_augmented} aug),"
+            f" {settings.sim.n_controls} controls, {settings.scp.n} nodes"
+        ),
+        f"Constraints: {n_nodal_convex} conv, {n_nodal_nonconvex} nonconv, {n_ctcs} ctcs",
+        (
+            f"Subproblem: {n_cvx_variables} vars, {n_cvx_parameters} params,"
+            f" {n_cvx_constraints} constraints"
+        ),
+        f"Weights: {weights_str}",
+        f"CVX Solver: {settings.cvx.solver}, Discretization Solver: {settings.dis.solver}",
+        f"JAX Backend: {jax_backend} (v{jax_version})",
+    ]
+
+    print_summary_box(lines, "Problem Summary")
+
+
+def print_results_summary(result: OptimizationResults, timing_post, timing_init, timing_solve):
+    """
+    Print the results summary box.
+
+    Args:
+        result (OptimizationResults): Optimization results object
+        timing_post (float): Post-processing time
+        timing_init (float): Initialization time
+        timing_solve (float): Solve time
+    """
+    cost = result.get("cost", 0.0)
+    ctcs_violation = result.get("ctcs_violation", 0.0)
+
+    # Convert numpy arrays to scalars for formatting
+    if hasattr(cost, "item"):
+        cost = cost.item()
+
+    # Handle CTCS violation - display as 1D array
+    if hasattr(ctcs_violation, "size"):
+        if ctcs_violation.size == 1:
+            ctcs_violation_str = f"[{ctcs_violation.item():.2e}]"
+        else:
+            # Display as 1D array
+            ctcs_violation_str = f"[{', '.join([f'{v:.2e}' for v in ctcs_violation])}]"
+    else:
+        ctcs_violation_str = f"[{ctcs_violation:.2e}]"
+
+    # Calculate total computation time
+    total_time = (timing_init or 0.0) + (timing_solve or 0.0) + timing_post
+
+    lines = [
+        "Results Summary",
+        f"Cost: {cost:.6f}",
+        f"CTCS Constraint Violation: {ctcs_violation_str}",
+        f"Preprocessing Time: {timing_init or 0.0:.3f}s",
+        f"Main Solve Time: {timing_solve or 0.0:.3f}s",
+        f"Post-processing Time: {timing_post:.3f}s",
+        f"Total Computation Time: {total_time:.3f}s",
+    ]
+
+    print_summary_box(lines, "Results Summary")
+
+
+def intro():
+    # Silence syntax warnings
+    warnings.filterwarnings("ignore")
+    # fmt: off
+    ascii_art = rf"""
+
+                            ____                    _____  _____
+                           / __ \                  / ____|/ ____|
+                          | |  | |_ __   ___ _ __ | (___ | |  __   ____  __
+                          | |  | | '_ \ / _ \ '_ \ \___ \| |  \ \ / /\ \/ /
+                          | |__| | |_) |  __/ | | |____) | |___\ V /  >  <
+                           \____/| .__/ \___|_| |_|_____/ \_____\_/  /_/\_\
+                                 | |
+                                 |_|
+─────────────────────────────────────────────────────────────────────────────────────────────────────────
+                                Author: Chris Hayner and Griffin Norris
+                                    Autonomous Controls Laboratory
+                                       University of Washington
+                                         Version: {get_version()}
+─────────────────────────────────────────────────────────────────────────────────────────────────────────
+"""
+    # fmt: on
+    print(ascii_art)
+
+
+def header():
+    print(
+        colored(
+            "─────────────────────────────────────────────────────────────────────────────────────────────────────────"
+        )
+    )
+    print(
+        "{:^4} │ {:^7} │ {:^7} │ {:^7} │ {:^7} │ {:^7} │ {:^7} │  {:^7} │ {:^14}".format(
+            "Iter",
+            "Dis Time (ms)",
+            "Solve Time (ms)",
+            "J_total",
+            "J_tr",
+            "J_vb",
+            "J_vc",
+            "Cost",
+            "Solver Status",
+        )
+    )
+    print(
+        colored(
+            "─────────────────────────────────────────────────────────────────────────────────────────────────────────"
+        )
+    )
+
+
+def intermediate(print_queue, params):
+    hz = 30.0
+    while True:
+        t_start = time.time()
+        try:
+            data = print_queue.get(timeout=1.0 / hz)
+            # remove bottom labels and line
+            if data["iter"] != 1:
+                sys.stdout.write("\x1b[1A\x1b[2K\x1b[1A\x1b[2K")
+            if data["prob_stat"][3] == "f":
+                # Only show the first element of the string
+                data["prob_stat"] = data["prob_stat"][0]
+
+            iter_colored = colored("{:4d}".format(data["iter"]))
+            J_tot_colored = colored("{:.1e}".format(data["J_total"]))
+            J_tr_colored = colored(
+                "{:.1e}".format(data["J_tr"]),
+                col_pos if data["J_tr"] <= params.scp.ep_tr else col_neg,
+            )
+            J_vb_colored = colored(
+                "{:.1e}".format(data["J_vb"]),
+                col_pos if data["J_vb"] <= params.scp.ep_vb else col_neg,
+            )
+            J_vc_colored = colored(
+                "{:.1e}".format(data["J_vc"]),
+                col_pos if data["J_vc"] <= params.scp.ep_vc else col_neg,
+            )
+            cost_colored = colored("{:.1e}".format(data["cost"]))
+            prob_stat_colored = colored(
+                data["prob_stat"], col_pos if data["prob_stat"] == "optimal" else col_neg
+            )
+
+            print(
+                "{:^4} │     {:^6.2f}    │      {:^6.2F}     │ {:^7} │ {:^7} │ {:^7} │ {:^7} │ "
+                " {:^7} │ {:^14}".format(
+                    iter_colored,
+                    data["dis_time"],
+                    data["subprop_time"],
+                    J_tot_colored,
+                    J_tr_colored,
+                    J_vb_colored,
+                    J_vc_colored,
+                    cost_colored,
+                    prob_stat_colored,
+                )
+            )
+
+            print(
+                colored(
+                    "─────────────────────────────────────────────────────────────────────────────────────────────────────────"
+                )
+            )
+            print(
+                "{:^4} │ {:^7} │ {:^7} │ {:^7} │ {:^7} │ {:^7} │ {:^7} │  {:^7} │ {:^14}".format(
+                    "Iter",
+                    "Dis Time (ms)",
+                    "Solve Time (ms)",
+                    "J_total",
+                    "J_tr",
+                    "J_vb",
+                    "J_vc",
+                    "Cost",
+                    "Solver Status",
+                )
+            )
+        except queue.Empty:
+            pass
+        time.sleep(max(0.0, 1.0 / hz - (time.time() - t_start)))
+
+
+def footer():
+    print(
+        colored(
+            "─────────────────────────────────────────────────────────────────────────────────────────────────────────"
+        )
+    )

openscvx/utils/profiling.py

@@ -0,0 +1,37 @@
+from datetime import datetime
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    import cProfile
+
+
+def profiling_start(profiling_enabled: bool) -> "Optional[cProfile.Profile]":
+    """Start profiling if enabled.
+
+    Args:
+        profiling_enabled: Whether to enable profiling.
+
+    Returns:
+        Profile object if enabled, None otherwise.
+    """
+    if profiling_enabled:
+        import cProfile
+
+        pr = cProfile.Profile()
+        pr.enable()
+        return pr
+    return None
+
+
+def profiling_end(pr: "Optional[cProfile.Profile]", identifier: str):
+    """Stop profiling and save results with timestamp.
+
+    Args:
+        pr: Profile object from profiling_start, or None.
+        identifier: Identifier for the profiling session (e.g., "solve", "initialize").
+    """
+    if pr is not None:
+        pr.disable()
+        # Save results so it can be visualized with snakeviz
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        pr.dump_stats(f"profiling/{timestamp}_{identifier}.prof")