torchsky.map.solvers¶
Experimental surface
These APIs are under active development. Names, defaults, and result objects may change between minor releases; pin a version if you depend on them in production.
Matrix-free linear solvers and bound observation operators for discrete inverse mapmaking: harmonic and T/E/B prior operators, operator composition, diagonal preconditioners, and spherical-harmonic solver hooks. Experimental: signatures may change between minor releases.
This reference is generated from the package source. Names are grouped by
the module that defines them; import every name from the public
namespace torchsky.map.solvers unless the group text says otherwise.
Core¶
Solver hooks for field fitting and discrete inverse mapmaking.
The implementation is grouped into _core (operators, probes,
preconditioners), _state (state-layout operators and component
decomposition), _priors (harmonic/TEB prior operators), and _hooks
(survey hooks). The historical :mod:torchsky.map.solvers names are re-exported
here.
BoundObservationOperator
dataclass
¶
BoundObservationOperator(
operators: ObservationOperatorBundle,
target: TargetRenderSpec,
apply_operator: ApplyOperatorFn | None = None,
apply_adjoint: AdjointOperatorFn | None = None,
apply_noise: ApplyNoiseFn | None = None,
)
Target-bound observation operator shared by discrete and INR scene types.
estimate_normal_operator_diagonal ¶
estimate_normal_operator_diagonal(
data: Tensor,
*,
initial_probe: Tensor | None = None,
damping: float = 0.0,
num_probes: int = 16,
) -> Tensor
estimate_normal_operator_scale ¶
estimate_normal_operator_scale(
data: Tensor,
*,
initial_probe: Tensor | None = None,
iterations: int = 8,
) -> float
estimate_posterior_variance_diagonal ¶
estimate_posterior_variance_diagonal(
reference_state: Tensor,
*,
num_probes: int = 32,
seed: int = 0,
iterations: int = 32,
damping: float = 0.0,
preconditioner: PreconditionerFn | Tensor | None = None,
rtol: float = 1e-05,
atol: float = 0.0,
) -> Tensor
Estimate the inverse-curvature diagonal with signed Hutchinson probes.
The target is
diag((A^T N^-1 A + damping I)^-1); damping changes curvature and
does not give the sampling covariance of a regularized estimator.
Negative finite-probe estimates remain signed, without a positive floor.
Every solve must satisfy the actual residual criterion
|| (H + damping I) x - z || <= max(atol, rtol ||z||) or raise.
The declared normal system must be fixed, Hermitian positive definite
on the chosen state space, with any unmeasured gauges removed explicitly.
A finite set of checked probes does not certify global identifiability
or establish a bound on inverse error for an ill-conditioned system.
A posterior interpretation additionally requires a declared Gaussian
likelihood and prior; inverse curvature alone supplies neither.
normal_equation_probe ¶
normal_equation_probe(
scene: SceneRepresentation | Module | Tensor,
data: Tensor,
) -> NormalEquationProbe
normal_operator_action ¶
normal_operator_action(
scene: SceneRepresentation | Module | Tensor,
) -> Tensor
Apply A^T N^-1 A without materializing the normal matrix.
recommended_step_size ¶
recommended_step_size(
data: Tensor,
*,
initial_probe: Tensor | None = None,
iterations: int = 8,
safety: float = 0.9,
) -> float
solve_least_squares_lsqr ¶
solve_least_squares_lsqr(
data: Tensor,
*,
initial_solution: Tensor | None = None,
iterations: int = 100,
rtol: float = 1e-05,
atol: float = 0.0,
damping: float = 0.0,
projection: GaugeProjection[Tensor]
| Callable[[Tensor], Tensor]
| None = None,
) -> LinearSolveResult
solve_normal_equation_cg ¶
solve_normal_equation_cg(
data: Tensor,
*,
initial_solution: Tensor | None = None,
iterations: int = 16,
damping: float = 0.0,
preconditioner: Preconditioner[Tensor]
| PreconditionerFn
| Tensor
| None = None,
rtol: float = 1e-11,
atol: float = 0.0,
projection: GaugeProjection[Tensor]
| Callable[[Tensor], Tensor]
| None = None,
operator_is_hpd: bool = True,
) -> LinearSolveResult
solve_normal_system_cg ¶
solve_normal_system_cg(
rhs: Tensor,
*,
initial_solution: Tensor | None = None,
iterations: int = 16,
damping: float = 0.0,
preconditioner: Preconditioner[Tensor]
| PreconditionerFn
| Tensor
| None = None,
rtol: float = 1e-11,
atol: float = 0.0,
projection: GaugeProjection[Tensor]
| Callable[[Tensor], Tensor]
| None = None,
operator_is_hpd: bool = True,
) -> LinearSolveResult
Solve (A^T N^-1 A + damping I) x = rhs matrix-free.
solve_weighted_least_squares ¶
solve_weighted_least_squares(
data: Tensor,
*,
initial_solution: Tensor | None = None,
iterations: int = 16,
step_size: float | None = None,
scale_iterations: int = 8,
step_safety: float = 0.9,
) -> LinearSolveResult
BoundObservationOperatorSet
dataclass
¶
BoundObservationOperatorSet(
operators: tuple[BoundObservationOperator, ...],
exposure_chunk_size: int = 1,
)
Shared-scene operator assembled from multiple target-bound observations.
_residual_adjoint_and_norm ¶
_residual_adjoint_and_norm(
scene: SceneRepresentation | Module | Tensor,
data: Sequence[Tensor],
) -> tuple[Tensor, Tensor]
_residual_norm ¶
_residual_norm(
scene: SceneRepresentation | Module | Tensor,
data: Sequence[Tensor],
) -> Tensor
estimate_normal_operator_diagonal ¶
estimate_normal_operator_diagonal(
data: Sequence[Tensor],
*,
initial_probe: Tensor | None = None,
damping: float = 0.0,
num_probes: int = 16,
) -> Tensor
estimate_normal_operator_scale ¶
estimate_normal_operator_scale(
data: Sequence[Tensor],
*,
initial_probe: Tensor | None = None,
iterations: int = 8,
) -> float
estimate_posterior_variance_diagonal ¶
estimate_posterior_variance_diagonal(
reference_state: Tensor,
*,
num_probes: int = 32,
seed: int = 0,
iterations: int = 32,
damping: float = 0.0,
preconditioner: PreconditionerFn | Tensor | None = None,
rtol: float = 1e-05,
atol: float = 0.0,
) -> Tensor
Estimate inverse curvature with signed, residual-checked probes.
Uses the same contract as
BoundObservationOperator.estimate_posterior_variance_diagonal:
no positive floor, failure on an inadequate actual solve residual,
and no claim that damped inverse curvature is sampling covariance.
Checked finite probes do not certify global identifiability.
iter_apply_chunks ¶
iter_apply_chunks(
scene: SceneRepresentation | Module | Tensor,
) -> Iterator[tuple[Tensor, ...]]
Yield predictions in bounded, deterministic exposure-order chunks.
normal_equation_probe ¶
normal_equation_probe(
scene: SceneRepresentation | Module | Tensor,
data: Sequence[Tensor],
) -> MultiObservationNormalEquationProbe
normal_operator_action ¶
normal_operator_action(
scene: SceneRepresentation | Module | Tensor,
) -> Tensor
Apply A^T N^-1 A with chunk-bounded prediction intermediates.
recommended_step_size ¶
recommended_step_size(
data: Sequence[Tensor],
*,
initial_probe: Tensor | None = None,
iterations: int = 8,
safety: float = 0.9,
) -> float
solve_least_squares_lsqr ¶
solve_least_squares_lsqr(
data: Sequence[Tensor],
*,
initial_solution: Tensor | None = None,
iterations: int = 100,
rtol: float = 1e-05,
atol: float = 0.0,
damping: float = 0.0,
projection: GaugeProjection[Tensor]
| Callable[[Tensor], Tensor]
| None = None,
) -> LinearSolveResult
solve_normal_equation_cg ¶
solve_normal_equation_cg(
data: Sequence[Tensor],
*,
initial_solution: Tensor | None = None,
iterations: int = 16,
damping: float = 0.0,
preconditioner: Preconditioner[Tensor]
| PreconditionerFn
| Tensor
| None = None,
rtol: float = 1e-11,
atol: float = 0.0,
projection: GaugeProjection[Tensor]
| Callable[[Tensor], Tensor]
| None = None,
operator_is_hpd: bool = True,
) -> LinearSolveResult
solve_normal_system_cg ¶
solve_normal_system_cg(
rhs: Tensor,
*,
initial_solution: Tensor | None = None,
iterations: int = 16,
damping: float = 0.0,
preconditioner: Preconditioner[Tensor]
| PreconditionerFn
| Tensor
| None = None,
rtol: float = 1e-11,
atol: float = 0.0,
projection: GaugeProjection[Tensor]
| Callable[[Tensor], Tensor]
| None = None,
operator_is_hpd: bool = True,
) -> LinearSolveResult
Solve (A^T N^-1 A + damping I) x = rhs matrix-free.
solve_weighted_least_squares ¶
solve_weighted_least_squares(
data: Sequence[Tensor],
*,
initial_solution: Tensor | None = None,
iterations: int = 16,
step_size: float | None = None,
scale_iterations: int = 8,
step_safety: float = 0.9,
) -> LinearSolveResult
weighted_adjoint ¶
weighted_adjoint(tensors: Sequence[Tensor]) -> Tensor
Apply A^T N^-1 without materializing weighted exposures.
compose_bound_operators ¶
compose_bound_operators(
*parts: BoundObservationOperator
| BoundObservationOperatorSet,
exposure_chunk_size: int = 1,
) -> BoundObservationOperatorSet
Concatenate bound operators and operator sets into one set.
Data operators and prior operators are stacked so a single normal-equation
solve sees all of them; exposure_chunk_size bounds how many are applied
at once.
make_diagonal_preconditioner ¶
make_diagonal_preconditioner(
diagonal: Tensor,
) -> DiagonalPreconditioner
Wrap a diagonal as a Jacobi-style :class:DiagonalPreconditioner.
Foundation¶
Solver hooks for field fitting and discrete inverse mapmaking.
The implementation is grouped into _core (operators, probes,
preconditioners), _state (state-layout operators and component
decomposition), _priors (harmonic/TEB prior operators), and _hooks
(survey hooks). The historical :mod:torchsky.map.solvers names are re-exported
here.
SolverDiagnostics
dataclass
¶
SolverDiagnostics(
termination: TerminationReason,
converged: bool,
iterations: int,
initial_residual_norm: float,
residual_norm: float,
relative_residual_norm: float,
stopping_threshold: float,
operator_applications: int,
adjoint_applications: int = 0,
preconditioner_applications: int = 0,
preconditioner_identity: str = "identity",
projected: bool = False,
breakdown: SolverBreakdown | None = None,
normal_residual_norm: float | None = None,
)
SolverResult
dataclass
¶
SolverResult(
solution: StateT,
residual_norm_history: tuple[float, ...] = (),
diagnostics: SolverDiagnostics | None = None,
)
Bases: Generic[StateT]
residual_norm_history
class-attribute
instance-attribute
¶
residual_norm_history: tuple[float, ...] = ()
conjugate_gradient ¶
conjugate_gradient(
operator: Callable[[StateT], StateT],
rhs: StateT,
*,
operator_is_hpd: bool,
initial_solution: StateT | None = None,
max_iterations: int = 100,
rtol: float = 1e-05,
atol: float = 0.0,
preconditioner: Preconditioner[StateT]
| Callable[[StateT], StateT]
| StateTree
| None = None,
projection: GaugeProjection[StateT]
| Callable[[StateT], StateT]
| None = None,
record_residual_history: bool = True,
) -> SolverResult[StateT]
Solve a declared HPD system using projected preconditioned CG.
lsqr ¶
lsqr(
operator: Callable[[StateT], ImageT],
adjoint: Callable[[ImageT], StateT],
rhs: ImageT,
*,
initial_solution: StateT | None = None,
max_iterations: int = 100,
rtol: float = 1e-05,
atol: float = 0.0,
damping: float = 0.0,
projection: GaugeProjection[StateT]
| Callable[[StateT], StateT]
| None = None,
record_residual_history: bool = True,
) -> SolverResult[StateT]
Solve a rectangular least-squares problem with Golub-Kahan LSQR.
Hooks¶
Solver hooks for field fitting and discrete inverse mapmaking.
The implementation is grouped into _core (operators, probes,
preconditioners), _state (state-layout operators and component
decomposition), _priors (harmonic/TEB prior operators), and _hooks
(survey hooks). The historical :mod:torchsky.map.solvers names are re-exported
here.
SolverHooks
dataclass
¶
SolverHooks(
operators: ObservationOperatorBundle,
apply_operator: ApplyOperatorFn | None = None,
apply_adjoint: AdjointOperatorFn | None = None,
apply_noise: ApplyNoiseFn | None = None,
bind_target_fn: BindTargetFn | None = None,
)
Internal assembly surface shared by GLS and MAP solver experiments.
bind_targets ¶
bind_targets(
targets: Sequence[TargetRenderSpec],
*,
exposure_chunk_size: int = 1,
) -> BoundObservationOperatorSet
normal_equation_probe ¶
normal_equation_probe(
scene: SceneRepresentation,
target: TargetRenderSpec,
data: Tensor,
) -> NormalEquationProbe
Assemble A^T N^-1 A and A^T N^-1 y probes without solving.
build_spherical_harmonic_solver_hooks ¶
build_spherical_harmonic_solver_hooks(
*,
lmax: int,
mmax: int | None = None,
apply_noise: ApplyNoiseFn | None = None,
) -> SolverHooks
Build native image↔alm solver hooks for direct harmonic projection.
build_spherical_spin_harmonic_solver_hooks ¶
build_spherical_spin_harmonic_solver_hooks(
*,
spin: int,
lmax: int,
mmax: int | None = None,
apply_noise: ApplyNoiseFn | None = None,
) -> SolverHooks
Build native image↔(E,B) alm solver hooks for direct spin harmonic projection.
Priors¶
Solver hooks for field fitting and discrete inverse mapmaking.
The implementation is grouped into _core (operators, probes,
preconditioners), _state (state-layout operators and component
decomposition), _priors (harmonic/TEB prior operators), and _hooks
(survey hooks). The historical :mod:torchsky.map.solvers names are re-exported
here.
build_harmonic_cl_prior_operator ¶
build_harmonic_cl_prior_operator(
reference_state: Tensor,
*,
lmax: int,
cls: Tensor | tuple[Tensor, Tensor] | list[Tensor],
mmax: int | None = None,
beam_transfer: Tensor
| tuple[Tensor, Tensor]
| list[Tensor]
| None = None,
min_variance: float = 1e-08,
apply_noise: ApplyNoiseFn | None = None,
) -> BoundObservationOperator
Bind a Gaussian harmonic prior with a per-multipole angular power spectrum.
The returned operator whitens spherical-harmonic coefficients by
1/sqrt(C_l) (optionally times a beam transfer function), for use as a
regulariser in a joint solve. reference_state fixes the coefficient
layout and dtype; min_variance floors the per-mode variance.
build_harmonic_degree_prior_operator ¶
build_harmonic_degree_prior_operator(
reference_state: Tensor,
*,
lmax: int,
mmax: int | None = None,
power: float = 1.0,
ell_offset: float = 0.0,
amplitude: float = 1.0,
include_monopole: bool = True,
apply_noise: ApplyNoiseFn | None = None,
) -> BoundObservationOperator
Bind a power-law harmonic prior amplitude * (l + ell_offset) ** -power.
Whitens spherical-harmonic coefficients by the inverse square root of the
degree-dependent variance. reference_state fixes the coefficient layout
and dtype; include_monopole controls the l = 0 mode.
build_harmonic_dense_prior_operator ¶
build_harmonic_dense_prior_operator(
reference_state: Tensor,
covariance: Tensor,
*,
apply_noise: ApplyNoiseFn | None = None,
jitter: float = 1e-08,
) -> BoundObservationOperator
Bind a Gaussian harmonic prior from a dense coefficient covariance.
The precision (covariance + jitter * I) ** -1 is factorised once by
Cholesky and applied as a matrix operator on the flattened coefficients.
build_teb_cls_prior_operator ¶
build_teb_cls_prior_operator(
reference_state: Tensor,
*,
lmax: int,
cls_tt: Tensor,
cls_ee: Tensor,
cls_bb: Tensor,
cls_te: Tensor | None = None,
cls_tb: Tensor | None = None,
cls_eb: Tensor | None = None,
mmax: int | None = None,
beam_transfer_t: Tensor | None = None,
beam_transfer_pol: Tensor | None = None,
min_variance: float = 1e-08,
jitter: float = 1e-08,
apply_noise: ApplyNoiseFn | None = None,
) -> BoundObservationOperator
Bind a joint T/E/B Gaussian prior from auto and cross spectra.
reference_state has shape (3, nalm). The per-multipole 3x3 spectra
C_l (cls_tt, cls_ee, cls_bb and optional cross terms) are
converted to precision factors; jitter regularises singular blocks.
harmonic_dense_covariance ¶
harmonic_dense_covariance(
*,
lmax: int,
mmax: int | None = None,
ell_correlation_length: float = 1.0,
diagonal_jitter: float = 0.001,
amplitude: float = 1.0,
dtype: dtype = float64,
device: device | None = None,
) -> Tensor
Dense squared-exponential covariance over harmonic coefficient index.
Entries are amplitude * exp(-0.5 * (dl / ell_correlation_length) ** 2)
plus diagonal_jitter on the diagonal, for coefficients up to lmax
and mmax. Returns a (nalm, nalm) tensor.
State¶
Solver hooks for field fitting and discrete inverse mapmaking.
The implementation is grouped into _core (operators, probes,
preconditioners), _state (state-layout operators and component
decomposition), _priors (harmonic/TEB prior operators), and _hooks
(survey hooks). The historical :mod:torchsky.map.solvers names are re-exported
here.
build_state_identity_operator ¶
build_state_identity_operator(
reference_state: Tensor,
*,
apply_noise: ApplyNoiseFn | None = None,
) -> BoundObservationOperator
Bind the identity map on a flattened scene state.
Useful as a ridge or anchor term when composing with data operators via
:func:compose_bound_operators.
build_state_subspace_operator ¶
build_state_subspace_operator(
reference_state: Tensor,
base_operator: BoundObservationOperator
| BoundObservationOperatorSet,
selector: int | slice | Sequence[int] | Tensor,
) -> BoundObservationOperator | BoundObservationOperatorSet
Restrict base_operator to a subset of the leading state axis.
selector picks rows of the state (index, slice, sequence, or tensor);
the adjoint embeds results back into a zero state of the original shape.