Refactor for v1

.pre-commit-config.yaml

@@ -23,8 +23,10 @@ repos:
     hooks:
       - id: mypy
         exclude: examples|docs
+        args: [--always-false=USE_DOUBLE_PRECISION]
         additional_dependencies:
           - "qoolqit[extras]>=1.1" # The version in pyproject.toml must match
+          - "jaxtyping"
 
 # Cleanup jupyter notebooks
 -   repo: https://github.com/srstevenson/nb-clean

docs/content/classical/cplex_solver/cplex_solving.md

@@ -11,13 +11,13 @@ class CplexSolver(BaseClassicalSolver):
     def __init__(
         self,
         instance: QUBOInstance,
-        config: Optional[Dict[str, Any]] = None
+        config: Optional[dict[str, Any]] = None
     )
 ```
 
 - **Parameters**:
   - `instance` (`QUBOInstance`): QUBO problem containing a square coefficient matrix (`torch.Tensor`).
-  - `config` (`Optional[Dict[str, Any]]`): Dictionary supporting:
+  - `config` (`Optional[dict[str, Any]]`): Dictionary supporting:
     - `cplex_maxtime` (`float`, default `600.0`): Maximum solve time in seconds.
     - `cplex_log_path` (`str`, default `"solver.log"`): Path for CPLEX log output.
 

pyproject.toml

@@ -64,6 +64,8 @@ dev = [
   "jupyter",
   "tqdm",
   "mypy",
+  "jaxtyping",
+  "beartype",
 ]
 doc = [
   "mkdocs",

pyrightconfig.json

@@ -0,0 +1,5 @@
+{
+  "defineConstant": {
+    "USE_DOUBLE_PRECISION": false
+  }
+}

qubosolver/__init__.py

@@ -1,16 +1,65 @@
 from __future__ import annotations
 
-from importlib import import_module
+from qubosolver.types._checks import _RUNTIME_CHECKS
 
-from .data import *  # noqa: F403
-from .qubo_instance import *  # noqa: F403
-from .qubo_types import *  # noqa: F403
-from .config import *  # noqa: F403
-from .utils import *  # noqa: F403
+from qubosolver.types import (
+    bitstring,
+    bitstrings,
+    matrix,
+    tensor,
+    vector,
+    vectori,
+    Bitstring,
+    Bitstrings,
+    Matrix,
+    Tensor,
+    Vector,
+    Vectori,
+    QUBOSolution,
+    QUBOAnalyzer,
+    QUBOInstance,
+    EmbedderType,
+    DriveType,
+    LayoutType,
+    SolutionStatusType,
+    DensityType,
+)
+from qubosolver.config import SolverConfig, EmbeddingConfig, DriveShapingConfig
 
-list_of_submodules = [".data", ".utils", ".qubo_instance", ".config"]
+__all__ = [
+    # Submodules
+    "bitstring",
+    "bitstrings",
+    "matrix",
+    "tensor",
+    "vector",
+    "vectori",
+    # Type Aliases
+    "Bitstring",
+    "Bitstrings",
+    "Matrix",
+    "Tensor",
+    "Vector",
+    "Vectori",
+    # Classes
+    "QUBOSolution",
+    "QUBOAnalyzer",
+    "QUBOInstance",
+    # Enums
+    "EmbedderType",
+    "DriveType",
+    "LayoutType",
+    "SolutionStatusType",
+    "DensityType",
+    # Configs
+    "SolverConfig",
+    "EmbeddingConfig",
+    "DriveShapingConfig",
+    # Functions
+]
 
-__all__ = []
-for submodule in list_of_submodules:
-    __all_submodule__ = getattr(import_module(submodule, package="qubosolver"), "__all__")
-    __all__ += __all_submodule__
+if _RUNTIME_CHECKS:
+    from beartype import BeartypeConf
+    from beartype.claw import beartype_this_package
+
+    beartype_this_package(conf=BeartypeConf(warning_cls_on_decorator_exception=None))

qubosolver/_drive_shaping/__init__.py

@@ -0,0 +1,4 @@
+from __future__ import annotations
+
+from qubosolver.drive_shaping import heuristic as heuristic
+from qubosolver.drive_shaping import optimized as optimized

qubosolver/_drive_shaping/device_specs.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import math
+import warnings
+
+import qoolqit
+
+
+def compare_specs(a: dict, b: dict) -> None:
+    keys_a = set(a.keys())
+    keys_b = set(b.keys())
+
+    if keys_a != keys_b:
+        only_in_a = keys_a - keys_b
+        only_in_b = keys_b - keys_a
+        if only_in_a:
+            warnings.warn(f"Keys present in a but missing in b: {only_in_a}")
+        if only_in_b:
+            warnings.warn(f"Keys present in b but missing in a: {only_in_b}")
+
+    for key in keys_a & keys_b:
+        v1 = a[key]
+        v2 = b[key]
+        if v1 is None or v2 is None:
+            continue
+        if not math.isclose(float(v1), float(v2)):
+            warnings.warn(f"Value mismatch for key '{key}': a={v1}, b={v2}")
+
+
+def pulser_specs(
+    device: qoolqit.Device, normalize: bool = False, check_against_qoolqit: bool = False
+) -> dict[str, float | None]:
+    pulser_device = device._device
+    channel = pulser_device.channels["rydberg_global"]
+    specs: dict[str, float | None] = {}
+    specs["max_duration"] = (
+        float(pulser_device.max_sequence_duration)
+        if pulser_device.max_sequence_duration is not None
+        else None
+    )
+    specs["max_amplitude"] = channel.max_amp
+    specs["max_abs_detuning"] = channel.max_abs_detuning
+    specs["min_distance"] = pulser_device.min_atom_distance
+    specs["max_radial_distance"] = pulser_device.max_radial_distance
+    specs["min_avg_amp"] = channel.min_avg_amp
+    specs["dmm_bottom_detuning"] = None
+
+    dmm_channels = list(getattr(pulser_device, "dmm_channels", {}).values())
+    if dmm_channels:
+        specs["dmm_bottom_detuning"] = getattr(dmm_channels[0], "bottom_detuning", None)
+
+    if not normalize:
+        return specs
+
+    def _normalize(name: str, scale: float) -> None:
+        if specs[name] is not None:
+            specs[name] /= scale  # type: ignore[operator]
+
+    r0 = specs["min_distance"]
+    assert r0 is not None
+    _normalize("min_distance", r0)
+    _normalize("max_radial_distance", r0)
+
+    C6 = pulser_device.interaction_coeff
+    J0 = C6 / (r0**6)
+
+    _normalize("max_amplitude", J0)
+    _normalize("max_abs_detuning", J0)
+    _normalize("min_avg_amp", J0)
+    _normalize("dmm_bottom_detuning", J0)
+
+    # J0 is in rad/us and t in ns, hence the factor 1000
+    _normalize("max_duration", 1000.0 / J0)
+
+    if check_against_qoolqit:
+        qq_specs = qoolqit_specs(device)
+        compare_specs(specs, qq_specs)
+
+    return specs
+
+
+def qoolqit_specs(
+    device: qoolqit.Device, complete_with_pulser: bool = False, check_against_pulser: bool = False
+) -> dict[str, float | None]:
+    specs = device.specs
+    _pulser_specs = pulser_specs(device, normalize=True)
+
+    def import_from_pulser_or_set(name: str, fallback: float | None = None) -> None:
+        if name in specs.keys():
+            return
+        if complete_with_pulser:
+            specs[name] = _pulser_specs.get(name, fallback)
+        else:
+            specs[name] = fallback
+
+    # Don't merge, update specific keys that are known not be in Qoolqit specs
+    import_from_pulser_or_set("min_avg_amp")
+    import_from_pulser_or_set("dmm_bottom_detuning")
+
+    if check_against_pulser:
+        compare_specs(specs, _pulser_specs)
+
+    return specs