"""Fit bare chip parameters to a numerical dressed specification.
Devices and couplings declare how their input numbers map to dressed targets.
The fitter reads the desired chip structurally rather than diagonalizing it,
and a bounded nonlinear least-squares solve evaluates the candidate chips.
References
----------
Dispersive regime (cross-Kerr and dressed frequencies):
- Koch et al., *Phys. Rev. A* **76**, 042319 (2007),
"Charge-insensitive qubit design derived from the Cooper pair box"
— the DuffingTransmon approximation and its dispersive shifts.
- Gambetta et al., *Phys. Rev. A* **74**, 042318 (2006),
"Qubit-photon interactions in a cavity" — dispersive frequency shifts.
Coupling seeding here uses a numerical signed cross-Kerr root, requiring
a bracket across ``seed_strength_bounds`` without assuming monotonicity.
JAX traceability
----------------
Every bare parameter here (``device.freq``, ``device.anharmonicity``, the
``coupling_strength`` of a coupling) is a sweepable, differentiable quantity. A
chip that uses a JAX-native backend supplies a traced dressed-observable
residual and an exact Jacobian. SciPy uses their concrete values only at the
bounded trust-region boundary. The optimizer itself is not JAX-traceable, but
the output :class:`~quchip.chip.chip.Chip` stays fully traceable for every
downstream operation.
"""
from __future__ import annotations
import warnings
from collections.abc import Mapping
from numbers import Integral
from typing import Any, Literal
import jax
import jax.numpy as jnp
import numpy as np
from scipy.optimize import brentq, least_squares
from quchip.analysis.effective_hamiltonian import effective_hamiltonian_between_states
from quchip.backend import _backend_context
from quchip.chip import Chip
from quchip.cutoff import active_cutoff
from quchip.inverse_design.observables import (
TargetSpec,
build_dressed_target_specs,
)
from quchip.inverse_design.subsystems import (
build_local_subsystem,
device_labels_for_local_eval,
)
from quchip.inverse_design.types import FitADressResult, FitParameterReport, ObservableReport
from quchip.utils.labeling import resolve_label
# Largest total energy-level index that each target kind reads from the dressed spectrum.
_TARGET_LEVEL_SUMS = {"freq": 1, "exchange_rate": 1, "anharmonicity": 2, "cross_kerr": 2}
def _pair_labels(label: Any, kind: str) -> tuple[str, str]:
"""Validate a 2-tuple label for pair observables (``zz``, ``exchange``)."""
if not isinstance(label, tuple) or len(label) != 2:
raise ValueError(f"{kind} targets require a 2-tuple of device labels, got {label!r}")
return str(label[0]), str(label[1])
def _static_exchange_rate(chip: Chip, label: Any) -> Any:
"""Read the real single-excitation off-diagonal element in GHz through
the public effective_hamiltonian_between_states calculation."""
a, b = _pair_labels(label, "exchange")
n = len(chip.devices)
state_a = [0] * n
state_b = [0] * n
state_a[chip.device_index(a)] = 1
state_b[chip.device_index(b)] = 1
h_eff = effective_hamiltonian_between_states(chip, tuple(state_a), tuple(state_b))
return jnp.real(h_eff[0, 1])
def _estimate_bare_g(
chip: Chip,
coupling,
spec: TargetSpec,
seed_strength_bounds: tuple[float, float] = (1e-6, 0.25),
) -> float:
"""Root-seed the signed full cross-Kerr on an isolated two-device copy.
Require the target to be bracketed by seed_strength_bounds; monotonicity
is unnecessary for brentq. Preserve the coupling copy/rebind path, including
constructor-only state, and the parent backend. Set trials through the
coupling-strength interface. An unbracketed target raises, never saturates."""
target_val = float(spec.target)
if target_val == 0.0:
return 0.0
dev_a = coupling.device_a.copy()
dev_b = coupling.device_b.copy()
device_map = {dev_a.label: dev_a, dev_b.label: dev_b}
sub_coupling = coupling.copy(device_map)
sub = Chip(
[dev_a, dev_b],
[sub_coupling],
frame=chip.frame,
approximation=chip.approximation,
basis=chip.basis,
backend=chip.backend,
space=chip.space,
)
def obs_at_strength(strength: float) -> float:
sub_coupling.set_coupling_strength(strength)
return float(sub.static_zz(sub_coupling.device_a, sub_coupling.device_b))
lo, hi = seed_strength_bounds
obs_lo, obs_hi = obs_at_strength(lo), obs_at_strength(hi)
if not min(obs_lo, obs_hi) <= target_val <= max(obs_lo, obs_hi):
raise ValueError(
f"Cannot seed a bare coupling strength for {coupling.label!r}'s target "
f"{spec.kind}={target_val!r}: not bracketed by the observable at "
f"seed_strength_bounds={seed_strength_bounds!r} "
f"(observable(lo)={obs_lo!r}, observable(hi)={obs_hi!r})."
)
return brentq(lambda strength: obs_at_strength(strength) - target_val, lo, hi)
def _selected_parameter_names(
chip: Chip,
device_selection: Mapping[str, tuple[str, ...]],
coupling_selection: Mapping[str, tuple[str, ...]],
) -> list[str]:
"""Free parameter paths in chip order and declared parameter order."""
names = [
f"{device.label}.{name}"
for device in chip.devices
for name in device.tunable_params()
if name in device_selection.get(device.label, ())
]
names.extend(
f"{coupling.label}.{coupling.coupling_strength_name}"
for coupling in chip.couplings
if coupling.coupling_strength_name in coupling_selection.get(coupling.label, ())
)
return names
def _resolve_vary(
chip: Chip,
vary: Mapping,
) -> tuple[dict[str, tuple[str, ...]], dict[str, tuple[str, ...]]]:
"""Resolve complete device/coupling free-parameter allowlists.
Absent components and empty selections stay frozen. Accept objects or labels;
reject duplicate labels, unknown tunables and bare string selections. Return
(device_selection, coupling_selection) mappings."""
devices_by_label = {device.label: device for device in chip.devices}
couplings_by_label = {coupling.label: coupling for coupling in chip.couplings}
known_labels = sorted(set(devices_by_label) | set(couplings_by_label))
device_selection: dict[str, tuple[str, ...]] = {}
coupling_selection: dict[str, tuple[str, ...]] = {}
seen_labels: set[str] = set()
for key, value in vary.items():
label = resolve_label(key)
if label in seen_labels:
raise ValueError(f"vary has duplicate entries resolving to component {label!r}.")
seen_labels.add(label)
if isinstance(value, str):
raise ValueError(
f"vary[{label!r}] must be a collection of parameter names (e.g. ('{value}',)), "
f"got the bare string {value!r}."
)
try:
names = tuple(str(name) for name in value)
except TypeError as exc:
raise ValueError(f"vary[{label!r}] must be an iterable of parameter names, got {value!r}.") from exc
if label in devices_by_label:
device = devices_by_label[label]
available = tuple(device.tunable_params())
unknown = [name for name in names if name not in available]
if unknown:
raise ValueError(
f"vary[{label!r}] names {unknown} are not tunable parameters of "
f"{type(device).__name__} {label!r}. Available: {list(available)}."
)
device_selection[label] = names
elif label in couplings_by_label:
coupling = couplings_by_label[label]
available = (coupling.coupling_strength_name,)
unknown = [name for name in names if name not in available]
if unknown:
raise ValueError(
f"vary[{label!r}] names {unknown} are not the declared coupling-strength "
f"parameter of {type(coupling).__name__} {label!r}. Available: {list(available)}."
)
coupling_selection[label] = names
else:
raise ValueError(
f"vary key {label!r} does not match any device or coupling on the chip. "
f"Available labels: {known_labels}."
)
return device_selection, coupling_selection
def _pack_initial_params(
chip: Chip,
target_specs: tuple[TargetSpec, ...],
seed_strength_bounds: tuple[float, float] = (1e-6, 0.25),
*,
device_selection: Mapping[str, tuple[str, ...]],
coupling_selection: Mapping[str, tuple[str, ...]],
) -> tuple[list[str], np.ndarray]:
"""Pack selected parameters in declaration order, root-seeding nonlinear
cross-Kerr targets; other couplings retain their declared strengths."""
names = _selected_parameter_names(chip, device_selection, coupling_selection)
name_set = set(names)
values: list[float] = []
for device in chip.devices:
for param_name, param_value in device.tunable_params().items():
if f"{device.label}.{param_name}" in name_set:
values.append(float(param_value))
coupling_specs = {
spec.label: spec for spec in target_specs if isinstance(spec.label, str) and spec.kind == "cross_kerr"
}
for coupling in chip.couplings:
key = f"{coupling.label}.{coupling.coupling_strength_name}"
if key not in name_set:
continue
spec = coupling_specs.get(coupling.label)
if spec is not None:
if not coupling.reduces_to_crosskerr:
values.append(float(spec.target))
else:
values.append(_estimate_bare_g(chip, coupling, spec, seed_strength_bounds))
else:
values.append(float(coupling.coupling_strength))
return names, np.asarray(values, dtype=float)
def _rebuild_candidate(chip: Chip, names: list[str], values: Any) -> Chip:
"""Rebind the selected fit parameters on an isolated chip copy."""
return chip.with_params(dict(zip(names, values, strict=True)))
def _working_chip(candidate: Chip, label: Any, evaluator: str) -> Chip:
"""Pick the full chip or a one-hop local neighborhood for this target."""
if evaluator != "local":
return candidate
return build_local_subsystem(candidate, device_labels_for_local_eval(candidate, label))
def _coupling_by_label(chip: Chip, label: Any):
"""Return the coupling on *chip* with the given label."""
return next(c for c in chip.couplings if c.label == label)
def _target_model(candidate: Chip, spec: TargetSpec, evaluator: str, max_hilbert_dim: int) -> Chip:
"""Resolve and check the actual model before each spectral target evaluation."""
locator = spec.label
if isinstance(locator, str) and locator in candidate.coupling_map:
coupling = candidate.coupling_map[locator]
locator = (coupling.device_a_label, coupling.device_b_label)
working = _working_chip(candidate, locator, evaluator)
if spec.kind in _TARGET_LEVEL_SUMS:
policy = active_cutoff(working)
if policy is not None and _TARGET_LEVEL_SUMS[spec.kind] > policy.max_total:
raise ValueError(
f"The {spec.kind} target for {spec.label!r} reads labels with total excitation "
f"{_TARGET_LEVEL_SUMS[spec.kind]}, outside the excitation cutoff N <= {policy.max_total}. "
"Raise max_total."
)
dimension = working.analysis._dressed_dimension(_TARGET_LEVEL_SUMS[spec.kind])
if dimension > max_hilbert_dim:
raise ValueError(
f"The {evaluator} model for {spec.label!r} diagonalizes a matrix of dimension {dimension}, "
f"exceeding max_hilbert_dim={max_hilbert_dim}. Increase the limit to evaluate this model."
)
return working
def _evaluate_checked_spec(candidate: Chip, spec: TargetSpec, evaluator: str, max_hilbert_dim: int) -> Any:
"""Apply the matrix budget to the selected model before evaluating its observable."""
working = _target_model(candidate, spec, evaluator, max_hilbert_dim)
return _evaluate_spec(working, spec, "full")
def _evaluate_spec(candidate: Chip, spec: TargetSpec, evaluator: str) -> Any:
"""Compute the observable this ``spec`` anchors, on the appropriate working chip."""
match spec.kind:
case "freq":
return _working_chip(candidate, spec.label, evaluator).freq(spec.label)
case "anharmonicity":
return _working_chip(candidate, spec.label, evaluator).dressed_anharmonicity(spec.label)
case "cross_kerr":
if isinstance(spec.label, tuple):
working = _working_chip(candidate, spec.label, evaluator)
a, b = _pair_labels(spec.label, "cross_kerr")
return working.static_zz(a, b)
coupling = _coupling_by_label(candidate, spec.label)
working = _working_chip(
candidate,
(coupling.device_a_label, coupling.device_b_label),
evaluator,
)
return working.static_zz(coupling.device_a_label, coupling.device_b_label)
case "exchange_rate":
if not isinstance(spec.label, tuple):
coupling = _coupling_by_label(candidate, spec.label)
label = (coupling.device_a_label, coupling.device_b_label)
else:
label = spec.label
return _static_exchange_rate(_working_chip(candidate, label, evaluator), label)
case "coupling_strength":
return _coupling_by_label(candidate, spec.label).coupling_strength
case _:
raise ValueError(f"Unknown spec kind {spec.kind!r}")
def _auto_bounds(chip: Chip, names: list[str], x0: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
"""Use device-owned parameter bounds and sign-symmetric coupling-strength bounds."""
devices_by_label = {device.label: device for device in chip.devices}
couplings_by_label = {coupling.label: coupling for coupling in chip.couplings}
lower: list[float] = []
upper: list[float] = []
for name, value in zip(names, x0, strict=True):
owner_label, _, param_name = name.rpartition(".")
device = devices_by_label.get(owner_label)
if device is not None:
lo, hi = device.tunable_param_bounds(param_name, float(value))
lower.append(float(lo))
upper.append(float(hi))
continue
coupling = couplings_by_label.get(owner_label)
if coupling is not None and param_name == coupling.coupling_strength_name:
mag = max(0.25, 2.0 * abs(value)) if value != 0.0 else 0.25
lower.append(-mag)
upper.append(mag)
continue
raise ValueError(f"Unsupported latent parameter {name!r}")
return np.asarray(lower, dtype=float), np.asarray(upper, dtype=float)
def _parameter_seed_metadata(
chip: Chip,
names: list[str],
target_specs: tuple[TargetSpec, ...],
*,
user_starts: set[str],
) -> dict[str, tuple[str, str | None]]:
"""Describe where each optimizer start and coupling-sign branch came from."""
couplings = {coupling.label: coupling for coupling in chip.couplings}
coupling_specs = {
spec.label: spec for spec in target_specs if isinstance(spec.label, str) and spec.kind == "cross_kerr"
}
metadata: dict[str, tuple[str, str | None]] = {}
for name in names:
owner, _, _ = name.rpartition(".")
coupling = couplings.get(owner)
if name in user_starts:
metadata[name] = ("user start", "user supplied" if coupling is not None else None)
continue
if coupling is None:
metadata[name] = (
"component declaration",
None,
)
continue
spec = coupling_specs.get(owner)
if spec is not None and not coupling.reduces_to_crosskerr:
metadata[name] = ("target value", "target sign")
elif spec is not None:
metadata[name] = ("isolated-pair root solve", "positive convention")
else:
metadata[name] = (
"component declaration",
"declared sign",
)
return metadata
def _identifiability_receipt(
jacobian: Any,
names: list[str],
parameter_scale: np.ndarray,
) -> dict[str, Any]:
"""Analyze the final residual Jacobian in the solver's scaled coordinates."""
matrix = np.asarray(jacobian, dtype=float)
if matrix.ndim == 1:
matrix = matrix.reshape(1, -1)
scaled = matrix * np.asarray(parameter_scale, dtype=float)[None, :]
_, singular_values, right_vectors = np.linalg.svd(scaled, full_matrices=True)
largest = float(singular_values[0]) if singular_values.size else 0.0
tolerance = float(np.finfo(float).eps * max(scaled.shape, default=0) * largest)
rank = int(np.count_nonzero(singular_values > tolerance))
n_parameters = len(names)
rank_deficient = rank < n_parameters
if rank_deficient:
condition_number = float("inf")
elif singular_values.size:
condition_number = float(singular_values[0] / singular_values[n_parameters - 1])
else:
condition_number = float("inf")
weak_directions: list[dict[str, Any]] = []
for index in range(rank, n_parameters):
vector = right_vectors[index]
magnitude = float(np.max(np.abs(vector)))
relative = vector / magnitude if magnitude else vector
weak_directions.append(
{
"singular_value": (float(singular_values[index]) if index < singular_values.size else 0.0),
"relative_weights": {
name: float(weight)
for name, weight in zip(names, relative, strict=True)
if abs(float(weight)) >= 1e-8
},
}
)
return {
"jacobian_rank": rank,
"jacobian_shape": tuple(int(size) for size in scaled.shape),
"jacobian_rank_tolerance": tolerance,
"jacobian_singular_values": tuple(float(value) for value in singular_values),
"jacobian_condition_number": condition_number,
"rank_deficient": rank_deficient,
"weak_parameter_directions": tuple(weak_directions),
"jacobian_parameter_scale": dict(zip(names, (float(value) for value in parameter_scale), strict=True)),
}
def _coerce_mapping(value: Mapping | None, name: str) -> dict:
"""Return *value* as a plain dict, treating ``None`` as empty and rejecting non-``Mapping`` input."""
if value is None:
return {}
if isinstance(value, Mapping):
return dict(value)
raise TypeError(f"{name} must be None or a Mapping, got {type(value).__name__}")
def _jax_residual_functions(
chip: Chip,
names: list[str],
target_specs: tuple[TargetSpec, ...],
evaluator: str,
targets: np.ndarray,
scales: np.ndarray,
max_hilbert_dim: int = 10_000,
):
"""Return JAX residual functions for a JAX-backed chip, else ``None``."""
backend = chip.backend
if backend.array_module is not jnp:
return None
targets_jax = jnp.asarray(targets)
scales_jax = jnp.asarray(scales)
def residual(values):
with _backend_context(backend):
candidate = _rebuild_candidate(chip, names, values)
observed = jnp.stack([jnp.asarray(_evaluate_checked_spec(candidate, spec, evaluator, max_hilbert_dim))
for spec in target_specs])
return (observed - targets_jax) / scales_jax
return jax.jit(residual), jax.jit(jax.jacrev(residual))
[docs]
def fit_a_dress(
chip: Chip,
*,
constraints: Mapping | None = None,
vary: Mapping | None = None,
start: Mapping | None = None,
evaluator: Literal["full", "local"] = "full",
max_hilbert_dim: int = 10_000,
seed_strength_bounds: tuple[float, float] = (1e-6, 0.25),
max_nfev: int = 1000,
) -> FitADressResult:
"""Find bare parameters for a numerically specified dressed chip.
``fit_a_dress(desired)`` reads component-owned target declarations without
evaluating ``desired``. Common spectral devices interpret their declared
frequencies and anharmonicities as dressed targets. A ``Capacitive`` edge
between two non-computational modes targets the dressed ``exchange_rate``.
Other ``Capacitive`` edges and ``CrossKerr`` couplings target the full
``cross_kerr = E11 - E10 - E01 + E00``. The returned
:attr:`~quchip.inverse_design.types.FitADressResult.chip` is a fitted clone,
and ``desired`` is never mutated.
``constraints`` adds or replaces numerical observables. ``vary`` replaces
the component-owned selection of free parameters. ``start`` replaces the
selected starting values.
Parameters
----------
chip
Desired dressed-chip specification. Component declarations supply
numerical targets, and no dressed analysis runs on this object. A chip
declared for simulation is not a specification. Its bare coupling
strengths become ``cross_kerr`` or ``exchange_rate`` targets, and a
frequency-only ``vary`` cannot meet them. State such targets explicitly
with ``constraints``, and read the unmet-target line of ``summary()``.
constraints
Additional ``{component_or_pair: {observable: value_or_none}}``
constraints. The supported canonical observables are ``"freq"``,
``"anharmonicity"``, ``"cross_kerr"``, ``"exchange_rate"``, and
``"coupling_strength"``. The function also accepts the aliases
``"zz"``/``"static_zz"``, ``"exchange"``, and ``"g"``. An explicit
value replaces the same component default. ``None`` removes it. Device
pairs do not have to be direct coupling edges.
vary
Complete desired-chip allowlist of the bare parameters that can move:
``{component_or_label: name_collection}``. If omitted, each component's
conservative inverse-design policy applies. Components that are absent
from an explicit mapping are frozen.
start
Optional ``{"<component>.<parameter>": value}`` replacements for the
optimizer's selected starting values. A key that ``vary`` does not
select raises an error rather than being silently ignored.
evaluator
``"full"`` evaluates each target on the complete model. ``"local"``
explicitly selects one-hop neighborhoods and omits effects of devices
outside them. A neighborhood keeps the ports whose targets it holds.
Target reports record the selected evaluator. Unsupported local model
contributions raise, e.g. a feedline that links a kept port to a
discarded one. A ``PortNetwork`` with matrix-form ``scattering`` also
raises when a neighborhood keeps only some of its planes. With the
mapping form ``{(output, input): amplitude}``, each neighborhood keeps
the entries of its own planes. A chip with an active
:class:`~quchip.cutoff.ExcitationCutoff` raises, because a
neighborhood would spend the same budget on fewer devices.
max_hilbert_dim
Maximum dimension of each matrix that a target evaluation
diagonalizes. A chip whose static model conserves the total
excitation number diagonalizes only the excitation sectors that its
targets read. Under an excitation cutoff, each target diagonalizes
the cutoff classes that its labels occupy, and every candidate keeps
the same cutoff. Other chips diagonalize the full product space. A
larger matrix raises before seeding or optimization. Each rebound or
traced target checks the limit again if its conservation proof
changes.
No different approximation is selected. Increase the value to permit a larger model.
seed_strength_bounds
``(lo, hi)`` magnitude bounds for the seed root solve of the bare coupling
strength (:func:`_estimate_bare_g`). This solve serves ``cross_kerr`` coupling
targets. The values at both endpoints must bracket the target observable, or
seeding raises :class:`ValueError` rather than silently returning a saturated
endpoint.
max_nfev
Maximum number of residual evaluations for the SciPy
Trust-Region Reflective solver.
Returns
-------
FitADressResult
Fitted chip clone, loss, residual history, per-target
:class:`ObservableReport` tuples, packed parameters, and
solver metadata.
Raises
------
ValueError
Raised when ``seed_strength_bounds`` does not bracket a cross-Kerr seed, an
automatic plan is underdetermined, ``vary`` contains unknown, duplicate or
invalid parameter selections, or a resource ceiling is exceeded.
Warns
-----
UserWarning
An explicit ``vary`` plan is underdetermined by the target count or by
the final scaled-Jacobian rank. A nonconverged automatic fit also
returns its candidate with a warning when the final Jacobian does not
have full rank. Converged rank-deficient automatic fits raise an error.
Notes
-----
The fitter normalizes residuals by ``max(|target|, 1e-9)``, so every anchor
contributes on an equal *relative*-error basis. A coupling's
scalar-strength bounds are symmetric around zero, because the sign of a
capacitive-type coupling is physical and the fitter must not constrain it.
The solver's convergence tolerances (``ftol``/``xtol``/``gtol`` =
``1e-11``) are fixed fitter policy, not options. So is the ``x_scale``
floor (``1e-3``, applied per parameter as ``max(abs(x0), 1e-3)``).
**Identifiability.** The count of free parameters against residuals is necessary but
not sufficient. The fitter therefore computes the final SVD rank and condition
number from normalized residuals, in the same scaled parameter coordinates that the
solver uses. A custom :class:`~quchip.declarative.models.DeviceModel` with a
discovered ``tunable_param_names`` (the derived default, not an explicit
declaration) is not automatically fit-ready. An unbounded parameter still needs a
:meth:`~quchip.devices.base.BaseDevice.tunable_param_bounds` rule before the
optimizer can search it.
**JAX traceability.** When the chip uses a JAX-native backend, the complete
map from parameters to residuals and its exact Jacobian are JAX-traceable.
SciPy receives their concrete values for bounded trust-region control. The
optimizer itself is not differentiated. Other backends keep SciPy's
numerical Jacobian. In both cases, the returned chip stays fully traceable
and differentiable.
"""
if evaluator not in ("full", "local"):
raise ValueError(f"Unknown evaluator {evaluator!r}; choose 'full' or 'local'.")
if evaluator == "local":
# A one-hop neighborhood would spend the global excitation budget on fewer devices.
from quchip.engine.routes import require_workflow
require_workflow(chip, "fit_a_dress(evaluator='local')")
if isinstance(max_hilbert_dim, bool) or not isinstance(max_hilbert_dim, Integral) or max_hilbert_dim < 1:
raise ValueError("max_hilbert_dim must be a positive integer.")
constraint_mapping = _coerce_mapping(constraints, "constraints") if constraints is not None else None
target_specs = build_dressed_target_specs(chip, constraint_mapping)
for spec in target_specs:
_target_model(chip, spec, evaluator, max_hilbert_dim)
if vary is None:
targeted_labels = {spec.label for spec in target_specs if isinstance(spec.label, str)}
device_selection = {
device.label: device.default_fit_parameters()
for device in chip.devices
if device.label in targeted_labels and device.default_fit_parameters()
}
coupling_selection = {
coupling.label: (coupling.coupling_strength_name,)
for coupling in chip.couplings
if coupling.label in targeted_labels
and any(spec.label == coupling.label and spec.kind != "coupling_strength" for spec in target_specs)
}
else:
device_selection, coupling_selection = _resolve_vary(chip, _coerce_mapping(vary, "vary"))
n_free_parameters = len(_selected_parameter_names(chip, device_selection, coupling_selection))
n_target_residuals = len(target_specs)
automatic_desired_plan = vary is None
if n_free_parameters == 0:
raise ValueError(
"vary selects zero free parameters; fit_a_dress has nothing to optimize. "
"List at least one device parameter or coupling strength with a non-empty selection."
)
underdetermined_by_count = n_free_parameters > n_target_residuals
if underdetermined_by_count:
message = (
f"fit_a_dress has {n_free_parameters} free parameters but only {n_target_residuals} target "
f"residuals; the fit is underdetermined by count. Select fewer parameters with vary. "
"Target count alone does not guarantee identifiability."
)
if automatic_desired_plan:
raise ValueError(message)
warnings.warn(message, UserWarning, stacklevel=2)
names, x0 = _pack_initial_params(
chip,
target_specs,
seed_strength_bounds,
device_selection=device_selection,
coupling_selection=coupling_selection,
)
user_starts: set[str] = set()
if start is not None:
start_mapping = _coerce_mapping(start, "start")
user_starts = set(start_mapping)
unknown = sorted(set(start_mapping) - set(names))
if unknown:
raise ValueError(
f"start contains parameters that are not selected by vary: {unknown}. Selected parameters: {names}."
)
indices = {name: index for index, name in enumerate(names)}
x0 = x0.copy()
for name, value in start_mapping.items():
x0[indices[name]] = float(value)
targets_arr = np.asarray([spec.target for spec in target_specs], dtype=float)
scales = np.maximum(np.abs(targets_arr), 1e-9)
lower, upper = _auto_bounds(chip, names, x0)
parameter_scale = np.maximum(np.abs(x0), 1e-3)
seed_metadata = _parameter_seed_metadata(
chip,
names,
target_specs,
user_starts=user_starts,
)
def concrete_residuals(x: np.ndarray) -> np.ndarray:
candidate = _rebuild_candidate(chip, names, x)
values = np.asarray([_evaluate_checked_spec(candidate, spec, evaluator, max_hilbert_dim)
for spec in target_specs], dtype=float)
return (values - targets_arr) / scales
jacobian_mode = "finite-difference"
residuals = concrete_residuals
jacobian: Any = "2-point"
try:
jax_functions = _jax_residual_functions(
chip,
names,
target_specs,
evaluator,
targets_arr,
scales,
max_hilbert_dim,
)
except ImportError:
jax_functions = None
if jax_functions is not None:
residual_jax, jacobian_jax = jax_functions
def residuals(x: np.ndarray) -> np.ndarray:
return np.asarray(residual_jax(jnp.asarray(x)), dtype=float)
def jacobian(x: np.ndarray) -> np.ndarray:
return np.asarray(jacobian_jax(jnp.asarray(x)), dtype=float)
jacobian_mode = "jax"
untracked_residuals = residuals
r0 = untracked_residuals(x0)
history: list[float] = [float(r0 @ r0)]
history_points: list[np.ndarray] = [np.asarray(x0, dtype=float).copy()]
history_values: list[np.ndarray] = [np.array(r0, dtype=float)]
def tracked_residuals(x: np.ndarray) -> np.ndarray:
"""Evaluate residuals and retain one record per distinct parameter vector."""
values = untracked_residuals(x)
point = np.asarray(x, dtype=float)
if not np.array_equal(point, history_points[-1]):
history.append(float(values @ values))
history_points.append(point.copy())
history_values.append(np.array(values, dtype=float))
return values
result = least_squares(
tracked_residuals,
x0=x0,
jac=jacobian,
bounds=(lower, upper),
method="trf",
x_scale=parameter_scale,
ftol=1e-11,
xtol=1e-11,
gtol=1e-11,
max_nfev=max_nfev,
)
final_loss = float(2.0 * result.cost)
final_values = untracked_residuals(result.x)
if not np.array_equal(history_points[-1], np.asarray(result.x, dtype=float)):
history.append(final_loss)
history_points.append(np.asarray(result.x, dtype=float).copy())
history_values.append(np.array(final_values, dtype=float))
else:
history[-1] = final_loss
fitted_chip = _rebuild_candidate(chip, names, result.x)
identifiability = _identifiability_receipt(result.jac, names, parameter_scale)
if identifiability["rank_deficient"]:
weak_names = sorted(
{
name
for direction in identifiability["weak_parameter_directions"]
for name in direction["relative_weights"]
}
)
message = (
f"fit_a_dress final Jacobian rank {identifiability['jacobian_rank']} for "
f"{n_free_parameters} free parameters; the fitted bare chip is not locally identifiable. "
f"Weak parameter directions involve {weak_names}."
)
if automatic_desired_plan and result.status > 0:
raise ValueError(message)
warnings.warn(message, UserWarning, stacklevel=2)
# Reconstruct observed values at start/end by inverting the residual scaling.
initial_vec = r0 * scales + targets_arr
final_vec = final_values * scales + targets_arr
initial_reports = tuple(
ObservableReport(
spec.kind,
spec.label,
spec.target,
float(initial),
float(initial),
evaluator,
spec.source,
)
for spec, initial in zip(target_specs, initial_vec, strict=True)
)
final_reports = tuple(
ObservableReport(
spec.kind,
spec.label,
spec.target,
float(initial),
float(final),
evaluator,
spec.source,
)
for spec, initial, final in zip(target_specs, initial_vec, final_vec, strict=True)
)
parameter_reports = tuple(
FitParameterReport(
name=name,
initial=float(initial),
final=float(final),
lower_bound=float(lo),
upper_bound=float(hi),
seed_source=seed_metadata[name][0],
sign_choice=seed_metadata[name][1],
)
for name, initial, final, lo, hi in zip(
names,
x0,
result.x,
lower,
upper,
strict=True,
)
)
return FitADressResult(
chip=fitted_chip,
loss=final_loss,
history=np.asarray(history, dtype=float),
history_params=tuple(dict(zip(names, point, strict=True)) for point in history_points),
history_residuals=np.asarray(history_values) * scales,
initial_targets=initial_reports,
final_targets=final_reports,
initial_params=dict(zip(names, x0, strict=True)),
final_params=dict(zip(names, result.x, strict=True)),
parameter_reports=parameter_reports,
solver_info={
"method": "trf",
"status": int(result.status),
"message": result.message,
"nfev": int(result.nfev),
"jacobian": jacobian_mode,
"n_free_parameters": n_free_parameters,
"n_target_residuals": n_target_residuals,
"underdetermined_by_count": underdetermined_by_count,
"history_axis": "distinct residual evaluation",
"n_recorded_evaluations": len(history),
"input_contract": "desired-chip",
**identifiability,
},
)
