pqls 0.1.0__py3-none-any.whl

pqls/__init__.py

@@ -0,0 +1,61 @@
+"""
+PQLS: High-performance JAX-accelerated Lindblad steady-state solvers for quantum atomic systems.
+"""
+
+from .atoms import calculate_rabi_frequency, get_atom, get_ground_state
+from .lindblad import (
+    solve_custom_network,
+    solve_ladder_system,
+    solve_quantum_ladder,
+    solve_steady_state,
+)
+from .types import (
+    ArrayLike,
+    ArrayOrScalar,
+    AtomicLevel,
+    AtomName,
+    Driver,
+    QuantumLadder,
+    QuantumState,
+)
+from .units import (
+    GHz_to_eV,
+    au_to_coulomb_metre,
+    coulomb_metre_to_au,
+    electric_field_to_rabi_freq,
+    eV_to_GHz,
+    hz_to_rad_s,
+    rabi_freq_to_electric_field,
+    rad_s_to_hz,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    # Core Solvers
+    "solve_steady_state",
+    "solve_custom_network",
+    "solve_ladder_system",
+    "solve_quantum_ladder",
+    # Data Models & Types
+    "ArrayLike",
+    "ArrayOrScalar",
+    "AtomicLevel",
+    "QuantumState",
+    "QuantumLadder",
+    "Driver",
+    "AtomName",
+    # Atom Helpers
+    "get_atom",
+    "get_ground_state",
+    "calculate_rabi_frequency",
+    # Units & Conversions
+    "au_to_coulomb_metre",
+    "coulomb_metre_to_au",
+    "eV_to_GHz",
+    "GHz_to_eV",
+    "electric_field_to_rabi_freq",
+    "rabi_freq_to_electric_field",
+    "hz_to_rad_s",
+    "rad_s_to_hz",
+]

pqls/atoms.py

@@ -0,0 +1,109 @@
+"""
+Atoms management and caching for ARC physics parameters.
+"""
+
+from functools import lru_cache
+
+from arc import Cesium, Rubidium85, Rubidium87
+
+from .types import (
+    ArrayOrScalar,
+    AtomicLevel,
+    AtomName,
+    AtomType,
+    QuantumState,
+)
+from .units import au_to_coulomb_metre, electric_field_to_rabi_freq
+
+_ATOM_MAP = {
+    "rb85": Rubidium85,
+    "rb87": Rubidium87,
+    "cs": Cesium,
+}
+
+
+@lru_cache(maxsize=3)
+def get_atom(atom_name: AtomName) -> AtomType:
+    """Returns a cached instance of the requested atom from the ARC package.
+
+    Parameters
+    ----------
+    atom_name : AtomName
+        The common name of the atom (e.g., 'Rb85', 'Cs').
+
+    Returns
+    -------
+    AtomType
+        The instantiated atom object from the ARC library.
+
+    Raises
+    ------
+    ValueError
+        If the provided atom_name is not in the supported list.
+    """
+    atom_cls = _ATOM_MAP.get(atom_name.lower())
+    if atom_cls is None:
+        raise ValueError(f"Atom {atom_name} not supported.")
+    return atom_cls()
+
+
+def get_ground_state(atom_name: AtomName) -> AtomicLevel:
+    """Returns the ground state of the specified atom.
+
+    Parameters
+    ----------
+    atom_name : AtomName
+        The common name of the atom.
+
+    Returns
+    -------
+    AtomicLevel
+        An object representing the ``n``, ``l``, and ``j`` quantum numbers
+        of the ground state.
+    """
+    arc_atom = get_atom(atom_name)
+    return AtomicLevel(n=arc_atom.groundStateN, l=0, j=0.5)
+
+
+def calculate_rabi_frequency[T: ArrayOrScalar](
+    atom: AtomName,
+    lower: QuantumState,
+    upper: QuantumState,
+    e_field: T,
+    q: int | None = None,
+) -> T:
+    r"""Calculates angular Rabi frequency (``rad/s``) between two atomic states for a given electric field.
+
+    Parameters
+    ----------
+    atom : str, AtomName, or ARC Atom object
+        The target alkali atom.
+    lower : QuantumState
+        Initial / lower energy state.
+    upper : QuantumState
+        Target / upper energy state.
+    e_field : ArrayOrScalar
+        Electric field amplitude in ``V/m`` (scalar or array).
+    q : int, optional
+        Laser polarization transition (``-1`` for :math:`\sigma^-`, ``0`` for :math:`\pi`, ``+1`` for :math:`\sigma^+`).
+        If None, inferred automatically from :math:`\Delta m_j`.
+
+    Returns
+    -------
+    ArrayOrScalar
+        Angular Rabi frequency :math:`\Omega` in ``rad/s`` matching the shape of `e_field`.
+    """
+    arc_atom = get_atom(atom)
+
+    if q is None:
+        q = round(upper.mj - lower.mj)
+
+    dipole_au = abs(
+        arc_atom.getDipoleMatrixElement(
+            *lower.to_arc_args(),
+            *upper.to_arc_args(),
+            q=q,
+        )
+    )
+    dipole_cm = au_to_coulomb_metre(dipole_au)
+    return electric_field_to_rabi_freq(e_field, dipole_cm)

pqls/lindblad.py

@@ -0,0 +1,421 @@
+"""
+High-performance JAX-accelerated solvers for the Lindblad master equation steady-state.
+Provides both raw matrix-level APIs and high-level ladder system functions.
+"""
+
+from collections.abc import Sequence
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from .atoms import get_atom
+from .types import ArrayOrScalar, Driver, QuantumLadder
+from .units import au_to_coulomb_metre, electric_field_to_rabi_freq, hz_to_rad_s
+
+jax.config.update("jax_enable_x64", True)
+
+
+# --- Private JAX Kernels ---
+
+
+def _build_hamiltonian(
+    detunings: jnp.ndarray,
+    rabis: jnp.ndarray,
+    from_idx: jnp.ndarray,
+    to_idx: jnp.ndarray,
+    n_levels: int,
+) -> jnp.ndarray:
+    """Constructs the N x N Hamiltonian matrix."""
+    H = jnp.zeros((n_levels, n_levels), dtype=jnp.complex128)
+
+    H = H.at[jnp.diag_indices(n_levels)].set(detunings)
+
+    H = H.at[(from_idx, to_idx)].set(rabis / 2.0)
+    H = H.at[(to_idx, from_idx)].set(jnp.conj(rabis) / 2.0)
+
+    return H
+
+
+def _build_collapse_operators(
+    decay_from: jnp.ndarray, decay_to: jnp.ndarray, gammas: jnp.ndarray, n_levels: int
+) -> jnp.ndarray:
+    """Constructs a stack of K sparse decay/jump operators."""
+    K = gammas.shape[0]
+    C_ops = jnp.zeros((K, n_levels, n_levels), dtype=jnp.complex128)
+
+    # C_k = sqrt(gamma) * |to><from|
+    sqrt_g = jnp.sqrt(jnp.maximum(gammas, 0.0))
+    C_ops = C_ops.at[jnp.arange(K), decay_to, decay_from].set(sqrt_g)
+
+    return C_ops
+
+
+def _extract_output(
+    rho_ss: jnp.ndarray,
+    transitions: list[tuple[int, int]],
+    coherence_index: int | None,
+    is_scalar: bool,
+) -> ArrayOrScalar:
+    """Extracts either target coherence or full density matrix and formats dimensions."""
+    if coherence_index is not None:
+        try:
+            from_s, to_s = transitions[coherence_index]
+        except IndexError as err:
+            raise IndexError(
+                f"coherence_index {coherence_index} out of range for {len(transitions)} transitions."
+            ) from err
+        res = np.asarray(jnp.imag(rho_ss[:, to_s, from_s]))
+        return float(res[0]) if is_scalar else res
+
+    res = np.asarray(rho_ss)
+    return res[0] if is_scalar else res
+
+
+def _solve_steady_state_single(H: jnp.ndarray, C_ops: jnp.ndarray) -> jnp.ndarray:
+    r"""Solves the steady-state Lindblad master equation for an unbatched system.
+
+    Parameters
+    ----------
+    H : jnp.ndarray
+        Hamiltonian matrix of shape ``(N, N)``.
+    C_ops : jnp.ndarray
+        Stack of collapse/jump operators of shape ``(K, N, N)``.
+
+    Returns
+    -------
+    jnp.ndarray
+        Steady-state density matrix :math:`\rho_{ss}` of shape ``(N, N)`` with unit trace.
+    """
+
+    N = H.shape[0]
+    Ident = jnp.eye(N, dtype=jnp.complex128)
+
+    L = -1j * (jnp.kron(H, Ident) - jnp.kron(Ident, H.T))
+
+    def make_decay_superop(C):
+        CTC = C.T.conj() @ C
+        return (
+            jnp.kron(C, C.conj())
+            - 0.5 * jnp.kron(CTC, Ident)
+            - 0.5 * jnp.kron(Ident, CTC.T)
+        )
+
+    if C_ops.shape[0] > 0:
+        L_decay = jnp.sum(jax.vmap(make_decay_superop)(C_ops), axis=0)
+        L = L + L_decay
+
+    L = L.at[0, :].set(0.0)
+    diag_indices = jnp.arange(0, N**2, N + 1)
+    L = L.at[0, diag_indices].set(1.0)
+
+    b = jnp.zeros(N**2, dtype=jnp.complex128).at[0].set(1.0)
+    rho_vec = jnp.linalg.solve(L, b)
+
+    return rho_vec.reshape((N, N))
+
+
+def _single_point_network(
+    d_arr, r_arr, from_idx, to_idx, C_ops, n_levels
+) -> jnp.ndarray:
+    """Maps parameters to a Hamiltonian and solves the steady state."""
+    H = _build_hamiltonian(d_arr, r_arr, from_idx, to_idx, n_levels)
+    return _solve_steady_state_single(H, C_ops)
+
+
+def _single_point_ladder(det_vec, rab_vec, gammas_vec, n_levels) -> jnp.ndarray:
+    """N-level ladder point solver composing shared JAX primitives."""
+    cum_dets = jnp.concatenate([jnp.zeros(1, dtype=jnp.float64), jnp.cumsum(det_vec)])
+    diagonals = -1.0 * cum_dets
+
+    from_idx = jnp.arange(n_levels - 1)
+    to_idx = jnp.arange(1, n_levels)
+
+    H = _build_hamiltonian(diagonals, rab_vec, from_idx, to_idx, n_levels)
+    C_ops = _build_collapse_operators(to_idx, from_idx, gammas_vec[1:], n_levels)
+
+    return _solve_steady_state_single(H, C_ops)
+
+
+_compiled_solve_single = jax.jit(_solve_steady_state_single)
+_compiled_solve_batch = jax.jit(jax.vmap(_solve_steady_state_single, in_axes=(0, None)))
+_compiled_batched_network = jax.jit(
+    jax.vmap(_single_point_network, in_axes=(0, 0, None, None, None, None)),
+    static_argnums=(5,),  # n_levels must be a static integer for matrix allocation
+)
+_compiled_batched_ladder = jax.jit(
+    jax.vmap(_single_point_ladder, in_axes=(0, 0, None, None)), static_argnums=(3,)
+)
+
+# --- Public APIs ---
+
+
+def solve_steady_state(hamiltonian: np.ndarray, c_ops: np.ndarray) -> np.ndarray:
+    r"""Solves for the steady state given the Hamiltonian and collapse operators.
+
+    Parameters
+    ----------
+    hamiltonian : np.ndarray
+        Hamiltonian matrix of shape ``(N, N)`` or batched ``(B, N, N)``.
+    c_ops : np.ndarray
+        Stack of collapse operators of shape ``(K, N, N)``. These are shared across the batch.
+
+    Returns
+    -------
+    np.ndarray
+        Steady-state density matrix :math:`\rho_{ss}` of shape ``(N, N)`` or ``(B, N, N)``.
+    """
+    if hamiltonian.ndim == 2:
+        # Single unbatched matrix solve
+        res = _compiled_solve_single(jnp.asarray(hamiltonian), jnp.asarray(c_ops))
+    else:
+        # Batched across axis 0
+        res = _compiled_solve_batch(jnp.asarray(hamiltonian), jnp.asarray(c_ops))
+
+    return np.asarray(res)
+
+
+def solve_custom_network(
+    transitions: list[tuple[int, int]],
+    decays: list[tuple[int, int, float]],
+    detunings: list[ArrayOrScalar],
+    rabis: list[ArrayOrScalar],
+    n_levels: int,
+    coherence_index: int | None = None,
+) -> ArrayOrScalar:
+    r"""Solves the steady state for an arbitrary transition network.
+
+    Parameters
+    ----------
+    transitions : list of tuple of (int, int)
+        List of coupled transitions represented as ``(from_state, to_state)`` indices.
+    decays : list of tuple of (int, int, float)
+        List of decay channels represented as ``(from_state, to_state, gamma)``.
+    detunings : list of ArrayOrScalar
+        List of detunings for each transition in ``rad/s``.
+        Each scalar will be held constant throughout the batch.
+    rabis : list of ArrayOrScalar
+        List of angular Rabi frequencies for each transition in ``rad/s``.
+        Each scalar will be held constant throughout the batch.
+    n_levels : int
+        Total number of levels :math:`N` in the system.
+    coherence_index : int, optional
+        If specified (e.g., ``0``, ``1``), returns :math:`\text{Im}(\rho)` for ``transitions[coherence_index]``.
+        If ``None``, returns the full density matrix batch of shape ``(Batch, N, N)``.
+
+    Returns
+    -------
+    ArrayOrScalar
+        Extracted coherence or full steady-state density matrix.
+    """
+
+    num_dets = len(detunings)
+    all_bcast = np.broadcast_arrays(*detunings, *rabis)
+    is_scalar = all_bcast[0].ndim == 0
+
+    if is_scalar:
+        all_bcast = [np.expand_dims(x, 0) for x in all_bcast]
+
+    dets_stacked = np.stack(all_bcast[:num_dets], axis=-1)
+    rabs_stacked = np.stack(all_bcast[num_dets:], axis=-1)
+
+    from_idx = jnp.array([t[0] for t in transitions])
+    to_idx = jnp.array([t[1] for t in transitions])
+
+    decay_from = jnp.array([d[0] for d in decays], dtype=jnp.int32)
+    decay_to = jnp.array([d[1] for d in decays], dtype=jnp.int32)
+    gamma_vals = jnp.array([d[2] for d in decays], dtype=jnp.float64)
+
+    # Pre-build collapse operators (these do not change during parameter sweeps)
+    C_ops = _build_collapse_operators(decay_from, decay_to, gamma_vals, n_levels)
+
+    rho_ss = _compiled_batched_network(
+        jnp.asarray(dets_stacked),
+        jnp.asarray(rabs_stacked),
+        from_idx,
+        to_idx,
+        C_ops,
+        n_levels,
+    )
+
+    return _extract_output(rho_ss, transitions, coherence_index, is_scalar)
+
+
+def solve_ladder_system(
+    detunings: list[ArrayOrScalar],
+    rabis: list[ArrayOrScalar],
+    gammas: list[float],
+    coherence_index: int | None = 0,
+) -> ArrayOrScalar:
+    r"""Solves the steady state for an :math:`N`-level ladder system.
+
+    Parameters
+    ----------
+    detunings : list of ArrayOrScalar
+        List of detunings for each adjacent transition in ``rad/s``.
+        Each scalar will be held constant throughout the batch.
+    rabis : list of ArrayOrScalar
+        List of angular Rabi frequencies for each adjacent transition in ``rad/s``.
+        Each scalar will be held constant throughout the batch.
+    gammas : list of float
+        List of decay rates :math:`\Gamma` for each level in ``rad/s``. Note that ``gammas[0]`` is typically ``0.0``.
+    coherence_index : int, optional
+        Index of the transition to extract the coherence :math:`\text{Im}(\rho)` from. Defaults to ``0``.
+        If ``None``, returns the full density matrix.
+
+    Returns
+    -------
+    ArrayOrScalar
+        Extracted coherence or full steady-state density matrix.
+    """
+    N = len(rabis) + 1
+    num_dets = len(detunings)
+    all_bcast = np.broadcast_arrays(*detunings, *rabis)
+    is_scalar = all_bcast[0].ndim == 0
+
+    if is_scalar:
+        all_bcast = [np.expand_dims(x, 0) for x in all_bcast]
+
+    dets_stacked = jnp.asarray(
+        np.stack(all_bcast[:num_dets], axis=-1), dtype=jnp.float64
+    )
+    rabs_stacked = jnp.asarray(
+        np.stack(all_bcast[num_dets:], axis=-1), dtype=jnp.complex128
+    )
+    gammas_vec = jnp.asarray(gammas, dtype=jnp.float64)
+
+    rho_ss = _compiled_batched_ladder(dets_stacked, rabs_stacked, gammas_vec, N)
+    transitions = [(i, i + 1) for i in range(N - 1)]
+
+    return _extract_output(rho_ss, transitions, coherence_index, is_scalar)
+
+
+def solve_quantum_ladder(
+    ladder: QuantumLadder,
+    drivers: Sequence[Driver],
+    detunings: Sequence[ArrayOrScalar | None] | None = None,
+    frequencies: Sequence[ArrayOrScalar | None] | None = None,
+    amplitudes: Sequence[ArrayOrScalar | None] | None = None,
+    temperature: float = 300.0,
+    max_n: int = 80,
+    coherence_index: int | None = 0,
+) -> ArrayOrScalar:
+    r"""Solves the steady-state for a physical atomic ladder system using ARC and JAX.
+
+    Parameters
+    ----------
+    ladder : QuantumLadder
+        The quantum ladder system defining the atomic levels and transitions.
+    drivers : Sequence of Driver
+        Sequence of electromagnetic drivers for each transition.
+    detunings : Sequence of ArrayOrScalar or None, optional
+        Overriding detunings in ``Hz``. If provided, overrides driver settings.
+        Each scalar will be held constant throughout the batch.
+    frequencies : Sequence of ArrayOrScalar or None, optional
+        Overriding drive frequencies :math:`f` in ``Hz``. If provided, overrides driver settings.
+        Each scalar will be held constant throughout the batch.
+    amplitudes : Sequence of ArrayOrScalar or None, optional
+        Overriding electric field amplitudes :math:`E` in ``V/m``. If provided, overrides driver settings.
+        Each scalar will be held constant throughout the batch.
+    temperature : float, default=300.0
+        Environmental temperature in Kelvin for blackbody decay calculations.
+    max_n : int, default=80
+        Maximum principal quantum number :math:`n` to consider for decay pathways in ARC.
+    coherence_index : int, optional
+        Index of the transition to extract the coherence :math:`\text{Im}(\rho)` from. Defaults to ``0``.
+        If ``None``, returns the full density matrix.
+
+    Returns
+    -------
+    ArrayOrScalar
+        Extracted coherence or full steady-state density matrix.
+
+    Raises
+    ------
+    ValueError
+        If the number of drivers does not match the number of transitions in the ladder.
+    """
+    n_transitions = ladder.num_transitions
+    n_levels = ladder.num_levels
+
+    if len(drivers) != n_transitions:
+        raise ValueError(
+            f"Expected {n_transitions} drivers for a {n_levels}-level ladder, got {len(drivers)}."
+        )
+
+    atom = get_atom(ladder.atom)
+
+    # Compute gammas (decay rates in rad/s) from ARC
+    gammas = [0.0] * n_levels
+    for i in range(1, n_levels):
+        lifetime_s = atom.getStateLifetime(
+            *ladder.states[i].level.to_arc_args(),
+            temperature=temperature,
+            includeLevelsUpTo=max_n,
+        )
+        gammas[i] = 1.0 / lifetime_s if lifetime_s > 0 else 0.0
+
+    # Resolve Frequencies, Detunings, and Rabi Frequencies
+    detunings_list: list[ArrayOrScalar] = []
+    rabis_list: list[ArrayOrScalar] = []
+
+    for i in range(n_transitions):
+        lower = ladder.states[i]
+        upper = ladder.states[i + 1]
+        driver = drivers[i]
+
+        if detunings is not None and i < len(detunings) and detunings[i] is not None:
+            detuning_rad_s = hz_to_rad_s(detunings[i])
+
+        elif (
+            frequencies is not None
+            and i < len(frequencies)
+            and frequencies[i] is not None
+        ):
+            freq_val = frequencies[i]
+            if driver.account_for_detuning:
+                f_res_hz = atom.getTransitionFrequency(
+                    *lower.level.to_arc_args(),
+                    *upper.level.to_arc_args(),
+                )
+                detuning_rad_s = hz_to_rad_s(freq_val - f_res_hz)
+            else:
+                detuning_rad_s = (
+                    np.zeros_like(freq_val) if not np.isscalar(freq_val) else 0.0
+                )
+
+        else:
+            if driver.account_for_detuning:
+                f_res_hz = atom.getTransitionFrequency(
+                    *lower.level.to_arc_args(),
+                    *upper.level.to_arc_args(),
+                )
+                detuning_rad_s = hz_to_rad_s(driver.frequency - f_res_hz)
+            else:
+                detuning_rad_s = 0.0
+
+        detunings_list.append(detuning_rad_s)
+
+        if amplitudes is not None and i < len(amplitudes) and amplitudes[i] is not None:
+            amp_v_m = amplitudes[i]
+        else:
+            amp_v_m = driver.amplitude
+
+        # Calculate transition dipole moment via ARC (a.u. -> C*m)
+        dipole_au = atom.getDipoleMatrixElement(
+            *lower.to_arc_args(),
+            *upper.to_arc_args(),
+            q=round(upper.mj - lower.mj),
+        )
+        dipole_c_m = au_to_coulomb_metre(abs(dipole_au))
+
+        # Convert electric field to angular Rabi frequency (rad/s)
+        rabi_rad_s = electric_field_to_rabi_freq(amp_v_m, dipole_c_m)
+        rabis_list.append(rabi_rad_s)
+
+    return solve_ladder_system(
+        detunings=detunings_list,
+        rabis=rabis_list,
+        gammas=gammas,
+        coherence_index=coherence_index,
+    )