rheedium.recon¶

Surface reconstruction analysis and modeling utilities.

Differentiable inverse problems for RHEED data.

Extended Summary¶

This package is reserved for inverse-modeling code that infers structure, defects, and processing parameters from experimental RHEED data. Forward procedural models, preprocessing steps, and differentiable surface-construction utilities live in rheedium.procs.

Routine Listings¶

ReconstructionResult: Result container returned by reconstruction solvers.
weighted_image_residual(): Build a weighted least-squares residual field between two images.
weighted_mean_squared_error(): Compute a normalized weighted mean-squared error.
OrientationFitResult: Result container for orientation-distribution fitting.
orientation_loss(): Compute a masked image loss for an orientation distribution.
fit_orientation_weights(): Recover discrete orientation weights on a fixed candidate support.
compute_fisher_information(): Fisher information for fitted orientation-weight logits.
estimate_weight_uncertainty(): Propagate Fisher information to 1σ weight uncertainties.
gauss_newton_least_squares(): Gauss-Newton optimizer for arbitrary least-squares residuals.
adam_optimize(): Adam optimizer for arbitrary scalar objectives.
adagrad_optimize(): Adagrad optimizer for arbitrary scalar objectives.
gauss_newton_reconstruction(): Reconstruct parameters by least-squares image matching.
adam_reconstruction(): Reconstruct parameters by minimizing an image-matching loss with Adam.
adagrad_reconstruction(): Reconstruct parameters by minimizing an image-matching loss with Adagrad.

rheedium.recon.weighted_image_residual(simulated_image: Float[Array, 'H W'], experimental_image: Float[Array, 'H W'], weight_map: Float[Array, 'H W'] | None = None) → Float[Array, 'H W'][source]¶

Build a weighted least-squares residual field.

Extended Summary¶

The returned residual is designed for Gauss-Newton and other least-squares solvers. When a weight_map is provided it is interpreted as a non-negative per-pixel reliability weight:

\[r_{ij} = \sqrt{w_{ij}}\,(I^{\text{sim}}_{ij} - I^{\text{exp}}_{ij})\]

This convention ensures that minimizing \(\sum r_{ij}^2\) matches the standard weighted least-squares objective.

param simulated_image:: Simulated detector image.
type simulated_image:: Float[Array, 'H W']
param experimental_image:: Experimental target image.
type experimental_image:: Float[Array, 'H W']
param weight_map:: Non-negative pixel weights. Values of zero exclude pixels from the fit. Negative entries are clipped to zero.
type weight_map:: Optional[Float[Array, 'H W']], default: None
returns:: residual – Weighted residual image with the same shape as the inputs.
rtype:: Float[Array, 'H W']

rheedium.recon.weighted_mean_squared_error(simulated_image: Float[Array, 'H W'], experimental_image: Float[Array, 'H W'], weight_map: Float[Array, 'H W'] | None = None) → float | Float[Array, ''][source]¶

Compute a normalized weighted mean-squared error.

Extended Summary¶

Without weights this reduces to the ordinary mean-squared error. With weights it computes:

\[\mathrm{WMSE} = \frac{\sum_{ij} w_{ij}(I^{\text{sim}}_{ij} - I^{\text{exp}}_{ij})^2} {\max\left(\sum_{ij} w_{ij}, 10^{-12}\right)}\]

param simulated_image:: Simulated detector image.
type simulated_image:: Float[Array, 'H W']
param experimental_image:: Experimental target image.
type experimental_image:: Float[Array, 'H W']
param weight_map:: Non-negative pixel weights. Negative entries are clipped to zero before normalization.
type weight_map:: Optional[Float[Array, 'H W']], default: None
returns:: loss – Weighted mean-squared error.
rtype:: Union[float, Float[Array, '']]

class rheedium.recon.ReconstructionResult(params: Any, objective_history: Float[Array, 'N'], gradient_norm_history: Float[Array, 'N'], step_norm_history: Float[Array, 'N'], iterations: Int[Array, ''], converged: Bool[Array, ''])[source]¶

Bases: NamedTuple

Container for reconstruction outputs and optimization traces.

params¶

Final reconstructed parameter pytree.

Type:: Any

objective_history¶

Objective value after each accepted optimization step.

Type:: Float[Array, "N"]

gradient_norm_history¶

L2 norm of the gradient-like search direction at each iteration.

Type:: Float[Array, "N"]

step_norm_history¶

L2 norm of the parameter update applied at each iteration.

Type:: Float[Array, "N"]

iterations¶

Number of recorded optimization iterations.

Type:: Int[Array, ""]

converged¶

True when a convergence tolerance was met before exhausting the iteration budget.

Type:: Bool[Array, ""]

params: Any¶: Alias for field number 0

objective_history: Float[Array, 'N']¶: Alias for field number 1

gradient_norm_history: Float[Array, 'N']¶: Alias for field number 2

step_norm_history: Float[Array, 'N']¶: Alias for field number 3

iterations: Int[Array, '']¶: Alias for field number 4

converged: Bool[Array, '']¶: Alias for field number 5

rheedium.recon.adagrad_optimize(initial_params: Any, objective_fn: Callable[[Any], float | Float[Array, '']], learning_rate: float | Float[Array, ''] = 1e-1, epsilon: float | Float[Array, ''] = 1e-8, max_iterations: int = 500, tolerance: float | Float[Array, ''] = 1e-6) → ReconstructionResult[source]¶

Minimize a scalar objective with the Adagrad optimizer.

Parameters:

initial_params (Any) – Initial parameter pytree.
objective_fn (Callable[[Any], Union[float, Float[Array, '']]]) – Scalar objective to minimize.
learning_rate (Union[float, Float[Array, '']], default: 1e-1) – Adagrad base learning rate. Default: 1e-1
epsilon (Union[float, Float[Array, '']], default: 1e-8) – Denominator stabilizer. Default: 1e-8
max_iterations (int, default: 500) – Maximum number of iterations. Default: 500
tolerance (Union[float, Float[Array, '']], default: 1e-6) – Convergence threshold applied to the gradient norm and update norm. Default: 1e-6

Returns:

result – Final parameters plus optimization traces.

Return type:

ReconstructionResult

rheedium.recon.adagrad_reconstruction(initial_params: Any, forward_model: Callable[[Any], Float[Array, 'H W']], experimental_image: Float[Array, 'H W'], weight_map: Float[Array, 'H W'] | None = None, postprocess_fn: Callable[[Float[Array, 'H W']], Float[Array, 'H W']] | None = None, learning_rate: float | Float[Array, ''] = 1e-1, epsilon: float | Float[Array, ''] = 1e-8, max_iterations: int = 500, tolerance: float | Float[Array, ''] = 1e-6, loss_fn: Callable[[Float[Array, 'H W'], Float[Array, 'H W'], Float[Array, 'H W'] | None], float | Float[Array, '']] = weighted_mean_squared_error) → ReconstructionResult[source]¶

Reconstruct parameters by minimizing an image-matching loss.

Parameters:

initial_params (Any) – Initial parameter pytree passed to forward_model.
forward_model (Callable[[Any], Float[Array, 'H W']]) – Differentiable simulator that maps parameters to a detector image.
experimental_image (Float[Array, 'H W']) – Target detector image.
weight_map (Optional[Float[Array, 'H W']], default: None) – Optional non-negative per-pixel weights.
postprocess_fn (Optional[Callable[[Float[Array, 'H W']], Float[Array, 'H W']]], default: None) – Float[Array, “H W”]], optional Optional transform applied to each simulated image before loss evaluation.
learning_rate (Union[float, Float[Array, '']], default: 1e-1) – Adagrad base learning rate. Default: 1e-1
epsilon (Union[float, Float[Array, '']], default: 1e-8) – Denominator stabilizer. Default: 1e-8
max_iterations (int, default: 500) – Maximum number of iterations. Default: 500
tolerance (Union[float, Float[Array, '']], default: 1e-6) – Convergence threshold. Default: 1e-6
loss_fn (Callable[[Float[Array, 'H W'], Float[Array, 'H W'], Optional[Float[Array, 'H W']]], Union[float, Float[Array, '']]], default: weighted_mean_squared_error) – Image loss used for optimization. Default: weighted_mean_squared_error()

Returns:

result – Final parameters plus optimization traces.

Return type:

ReconstructionResult

rheedium.recon.adam_optimize(initial_params: Any, objective_fn: Callable[[Any], float | Float[Array, '']], learning_rate: float | Float[Array, ''] = 1e-2, beta1: float | Float[Array, ''] = 0.9, beta2: float | Float[Array, ''] = 0.999, epsilon: float | Float[Array, ''] = 1e-8, max_iterations: int = 250, tolerance: float | Float[Array, ''] = 1e-6) → ReconstructionResult[source]¶

Minimize a scalar objective with the Adam optimizer.

Parameters:

initial_params (Any) – Initial parameter pytree.
objective_fn (Callable[[Any], Union[float, Float[Array, '']]]) – Scalar objective to minimize.
learning_rate (Union[float, Float[Array, '']], default: 1e-2) – Adam learning rate. Default: 1e-2
beta1 (Union[float, Float[Array, '']], default: 0.9) – Exponential decay factor for the first moment. Default: 0.9
beta2 (Union[float, Float[Array, '']], default: 0.999) – Exponential decay factor for the second moment. Default: 0.999
epsilon (Union[float, Float[Array, '']], default: 1e-8) – Denominator stabilizer. Default: 1e-8
max_iterations (int, default: 250) – Maximum number of iterations. Default: 250
tolerance (Union[float, Float[Array, '']], default: 1e-6) – Convergence threshold applied to the gradient norm and update norm. Default: 1e-6

Returns:

result – Final parameters plus optimization traces.

Return type:

ReconstructionResult

rheedium.recon.adam_reconstruction(initial_params: Any, forward_model: Callable[[Any], Float[Array, 'H W']], experimental_image: Float[Array, 'H W'], weight_map: Float[Array, 'H W'] | None = None, postprocess_fn: Callable[[Float[Array, 'H W']], Float[Array, 'H W']] | None = None, learning_rate: float | Float[Array, ''] = 1e-2, beta1: float | Float[Array, ''] = 0.9, beta2: float | Float[Array, ''] = 0.999, epsilon: float | Float[Array, ''] = 1e-8, max_iterations: int = 250, tolerance: float | Float[Array, ''] = 1e-6, loss_fn: Callable[[Float[Array, 'H W'], Float[Array, 'H W'], Float[Array, 'H W'] | None], float | Float[Array, '']] = weighted_mean_squared_error) → ReconstructionResult[source]¶

Reconstruct parameters by minimizing an image-matching loss.

Parameters:

initial_params (Any) – Initial parameter pytree passed to forward_model.
forward_model (Callable[[Any], Float[Array, 'H W']]) – Differentiable simulator that maps parameters to a detector image.
experimental_image (Float[Array, 'H W']) – Target detector image.
weight_map (Optional[Float[Array, 'H W']], default: None) – Optional non-negative per-pixel weights.
postprocess_fn (Optional[Callable[[Float[Array, 'H W']], Float[Array, 'H W']]], default: None) – Float[Array, “H W”]], optional Optional transform applied to each simulated image before loss evaluation.
learning_rate (Union[float, Float[Array, '']], default: 1e-2) – Adam learning rate. Default: 1e-2
beta1 (Union[float, Float[Array, '']], default: 0.9) – First-moment decay factor. Default: 0.9
beta2 (Union[float, Float[Array, '']], default: 0.999) – Second-moment decay factor. Default: 0.999
epsilon (Union[float, Float[Array, '']], default: 1e-8) – Denominator stabilizer. Default: 1e-8
max_iterations (int, default: 250) – Maximum number of iterations. Default: 250
tolerance (Union[float, Float[Array, '']], default: 1e-6) – Convergence threshold. Default: 1e-6
loss_fn (Callable[[Float[Array, 'H W'], Float[Array, 'H W'], Optional[Float[Array, 'H W']]], Union[float, Float[Array, '']]], default: weighted_mean_squared_error) – Image loss used for optimization. Default: weighted_mean_squared_error()

Returns:

result – Final parameters plus optimization traces.

Return type:

ReconstructionResult

rheedium.recon.gauss_newton_least_squares(initial_params: Any, residual_fn: Callable[[Any], Array], damping: float | Float[Array, ''] = 1e-3, step_scale: float | Float[Array, ''] = 1.0, max_iterations: int = 25, tolerance: float | Float[Array, ''] = 1e-6) → ReconstructionResult[source]¶

Minimize a least-squares objective with Gauss-Newton iterations.

Extended Summary¶

This solver targets objectives of the form \(\min_\theta \|r(\theta)\|_2^2\), where residual_fn returns the residual vector or tensor. Parameters may be any JAX pytree; internally they are flattened with jax.flatten_util.ravel_pytree().

param initial_params:: Initial parameter pytree.
type initial_params:: Any
param residual_fn:: Residual function returning any array-shaped residual.
type residual_fn:: Callable[[Any], Array]
param damping:: Levenberg-style diagonal damping added to the normal matrix. Default: 1e-3
type damping:: Union[float, Float[Array, '']], default: 1e-3
param step_scale:: Scalar multiplier applied to the Gauss-Newton step. Default: 1.0
type step_scale:: Union[float, Float[Array, '']], default: 1.0
param max_iterations:: Maximum number of iterations. Default: 25
type max_iterations:: int, default: 25
param tolerance:: Convergence threshold applied to the gradient norm and update norm. Default: 1e-6
type tolerance:: Union[float, Float[Array, '']], default: 1e-6
returns:: result – Final parameters plus optimization traces.
rtype:: ReconstructionResult

rheedium.recon.gauss_newton_reconstruction(initial_params: Any, forward_model: Callable[[Any], Float[Array, 'H W']], experimental_image: Float[Array, 'H W'], weight_map: Float[Array, 'H W'] | None = None, postprocess_fn: Callable[[Float[Array, 'H W']], Float[Array, 'H W']] | None = None, damping: float | Float[Array, ''] = 1e-3, step_scale: float | Float[Array, ''] = 1.0, max_iterations: int = 25, tolerance: float | Float[Array, ''] = 1e-6) → ReconstructionResult[source]¶

Reconstruct parameters by least-squares image matching.

Parameters:

initial_params (Any) – Initial parameter pytree passed to forward_model.
forward_model (Callable[[Any], Float[Array, 'H W']]) – Differentiable simulator that maps parameters to a detector image.
experimental_image (Float[Array, 'H W']) – Target detector image.
weight_map (Optional[Float[Array, 'H W']], default: None) – Non-negative per-pixel weights for least-squares fitting.
postprocess_fn (Optional[Callable[[Float[Array, 'H W']], Float[Array, 'H W']]], default: None) – Float[Array, “H W”]], optional Optional transform applied to each simulated image before it is compared against experimental_image.
damping (Union[float, Float[Array, '']], default: 1e-3) – Diagonal damping for the Gauss-Newton normal equations. Default: 1e-3
step_scale (Union[float, Float[Array, '']], default: 1.0) – Scalar multiplier applied to each update. Default: 1.0
max_iterations (int, default: 25) – Maximum number of iterations. Default: 25
tolerance (Union[float, Float[Array, '']], default: 1e-6) – Convergence threshold. Default: 1e-6

Returns:

result – Final parameters plus optimization traces.

Return type:

ReconstructionResult

class rheedium.recon.OrientationFitResult(fitted_distribution: OrientationDistribution, final_loss: Float[Array, ''], loss_history: Float[Array, 'N_steps'], converged: bool, n_iterations: int, residual_pattern: Float[Array, 'H W'])[source]¶

Bases: NamedTuple

Results from orientation distribution fitting.

fitted_distribution¶

Recovered orientation distribution.

Type:: OrientationDistribution

final_loss¶

Final scalar loss value.

Type:: Float[Array, ""]

loss_history¶

Loss value after each optimizer step.

Type:: Float[Array, "N_steps"]

converged¶

Whether the optimizer met its stopping tolerance.

Type:: bool

n_iterations¶

Number of recorded optimization iterations.

Type:: int

residual_pattern¶

Difference I_observed - I_fitted for diagnostics.

Type:: float_jax_image

fitted_distribution: OrientationDistribution¶: Alias for field number 0

final_loss: Float[Array, '']¶: Alias for field number 1

loss_history: Float[Array, 'N_steps']¶: Alias for field number 2

converged: bool¶: Alias for field number 3

n_iterations: int¶: Alias for field number 4

residual_pattern: Float[Array, 'H W']¶: Alias for field number 5

rheedium.recon.compute_fisher_information(simulate_fn: Callable[[float | Float[Array, '']], Float[Array, 'H W']], distribution: OrientationDistribution, noise_variance: float | Float[Array, ''] = 1.0, mask: Float[Array, 'H W'] | None = None, normalize: bool = True, n_mosaic_points: int | Integer[Array, ''] = 7) → Float[Array, 'M M'][source]¶

Compute Fisher information for the discrete weight logits.

Description¶

This function treats the discrete orientation weights as a softmax of unconstrained logits, holds the orientation angles and mosaic width fixed, and computes the Gaussian-noise Fisher information matrix for those logits.

param simulate_fn:: Forward simulation function.
type simulate_fn:: Callable[[Union[float, Float[Array, '']]], Float[Array, 'H W']]
param distribution:: Distribution about which the local Fisher information is computed.
type distribution:: OrientationDistribution
param noise_variance:: Assumed per-pixel Gaussian noise variance. Default: 1.0
type noise_variance:: Union[float, Float[Array, '']], default: 1.0
param mask:: Optional detector mask. Default: all pixels.
type mask:: Optional[Float[Array, 'H W']], default: None
param normalize:: If True, compute Fisher information on normalized detector patterns. Default: True
type normalize:: bool, default: True
param n_mosaic_points:: Quadrature points for mosaic broadening. Default: 7
type n_mosaic_points:: Union[int, Integer[Array, '']], default: 7
returns:: fisher – Fisher information matrix for the M discrete weight logits.
rtype:: Float[Array, 'M M']

rheedium.recon.estimate_weight_uncertainty(result: OrientationFitResult, simulate_fn: Callable[[float | Float[Array, '']], Float[Array, 'H W']], noise_variance: float | Float[Array, ''] = 1.0, mask: Float[Array, 'H W'] | None = None, normalize: bool = True, n_mosaic_points: int | Integer[Array, ''] = 7) → Float[Array, 'M'][source]¶

Estimate 1σ uncertainties on the fitted discrete weights.

Description¶

Fisher information is computed in the unconstrained softmax-logit parameterization and then propagated to the physical weights with the Jacobian of the softmax map.

param result:: Fitting result from fit_orientation_weights().
type result:: OrientationFitResult
param simulate_fn:: Forward simulation function.
type simulate_fn:: Callable[[Union[float, Float[Array, '']]], Float[Array, 'H W']]
param noise_variance:: Assumed per-pixel Gaussian noise variance. Default: 1.0
type noise_variance:: Union[float, Float[Array, '']], default: 1.0
param mask:: Optional detector mask. Default: all pixels.
type mask:: Optional[Float[Array, 'H W']], default: None
param normalize:: If True, compute uncertainties on normalized detector patterns. Default: True
type normalize:: bool, default: True
param n_mosaic_points:: Quadrature points for mosaic broadening. Default: 7
type n_mosaic_points:: Union[int, Integer[Array, '']], default: 7
returns:: uncertainties – Approximate 1σ uncertainty for each discrete orientation weight.
rtype:: Float[Array, 'M']

Notes

The returned uncertainties are conditional on the fitted orientation angles and mosaic broadening being held fixed.

rheedium.recon.fit_orientation_weights(observed_pattern: Float[Array, 'H W'], simulate_fn: Callable[[float | Float[Array, '']], Float[Array, 'H W']], candidate_angles_deg: Float[Array, 'M'], mask: Float[Array, 'H W'] | None = None, initial_weights: Float[Array, 'M'] | None = None, learning_rate: float | Float[Array, ''] = 0.1, n_iterations: int | Integer[Array, ''] = 500, convergence_tol: float | Float[Array, ''] = 1e-6, regularization_strength: float | Float[Array, ''] = 1e-4, entropy_weight: float | Float[Array, ''] = 1e-3, n_mosaic_points: int | Integer[Array, ''] = 7, normalize: bool = True, verbose: bool = False) → OrientationFitResult[source]¶

Optimize orientation weights to match an observed pattern.

Description¶

Given a fixed candidate angle support, optimize the discrete orientation weights together with a shared Gaussian mosaic width. The internal optimization is carried out in an unconstrained space with a specialized Adam update on (weight_logits, mosaic_param) and a JIT-compiled jax.lax.scan() loop.

param observed_pattern:: Experimental RHEED pattern to match.
type observed_pattern:: Float[Array, 'H W']
param simulate_fn:: Forward simulation function mapping phi_deg to a detector pattern.
type simulate_fn:: Callable[[Union[float, Float[Array, '']]], Float[Array, 'H W']]
param candidate_angles_deg:: Candidate orientation angles in degrees.
type candidate_angles_deg:: Float[Array, 'M']
param mask:: Pixel mask for loss computation. Default: all pixels.
type mask:: Optional[Float[Array, 'H W']], default: None
param initial_weights:: Initial weight guess. Default: uniform weights.
type initial_weights:: Optional[Float[Array, 'M']], default: None
param learning_rate:: Adam learning rate. Default: 0.1
type learning_rate:: Union[float, Float[Array, '']], default: 0.1
param n_iterations:: Maximum optimization steps. Default: 500
type n_iterations:: Union[int, Integer[Array, '']], default: 500
param convergence_tol:: Stop when the optimizer update norm falls below this tolerance. Default: 1e-6
type convergence_tol:: Union[float, Float[Array, '']], default: 1e-6
param regularization_strength:: L2 regularization on weight deviations from uniform. Default: 1e-4
type regularization_strength:: Union[float, Float[Array, '']], default: 1e-4
param entropy_weight:: Entropy regularization. Default: 1e-3
type entropy_weight:: Union[float, Float[Array, '']], default: 1e-3
param n_mosaic_points:: Quadrature points for mosaic broadening. Default: 7
type n_mosaic_points:: Union[int, Integer[Array, '']], default: 7
param normalize:: If True, normalize both observed and simulated patterns before comparison. Default: True
type normalize:: bool, default: True
param verbose:: Print the final fitted weights and loss. Default: False
type verbose:: bool, default: False
returns:: result – Fitting result with recovered distribution and diagnostics.
rtype:: OrientationFitResult

rheedium.recon.orientation_loss(distribution: OrientationDistribution, simulate_fn: Callable[[float | Float[Array, '']], Float[Array, 'H W']], observed_pattern: Float[Array, 'H W'], mask: Float[Array, 'H W'] | None = None, normalize: bool = True, regularization_strength: float | Float[Array, ''] = 0.0, entropy_weight: float | Float[Array, ''] = 0.0, reference_weights: Float[Array, 'M'] | None = None, n_mosaic_points: int | Integer[Array, ''] = 7) → Float[Array, ''][source]¶

Compute loss between observed and simulated patterns.

Description¶

Computes a masked mean-squared error between an observed detector image and the incoherently averaged forward model implied by an OrientationDistribution.

param distribution:: Trial orientation distribution.
type distribution:: OrientationDistribution
param simulate_fn:: Forward simulation function mapping phi_deg to a detector pattern.
type simulate_fn:: Callable[[Union[float, Float[Array, '']]], Float[Array, 'H W']]
param observed_pattern:: Experimental or synthetic detector image to match.
type observed_pattern:: Float[Array, 'H W']
param mask:: Pixel mask for loss computation. Default: all pixels.
type mask:: Optional[Float[Array, 'H W']], default: None
param normalize:: If True, normalize both patterns before comparison. Default: True
type normalize:: bool, default: True
param regularization_strength:: L2 penalty on deviations from reference_weights. If reference_weights is omitted, a uniform distribution is used. Default: 0.0
type regularization_strength:: Union[float, Float[Array, '']], default: 0.0
param entropy_weight:: Weight on the entropy bonus -H(w). Positive values discourage collapse onto a single orientation. Default: 0.0
type entropy_weight:: Union[float, Float[Array, '']], default: 0.0
param reference_weights:: Reference weight vector for the L2 penalty. Default: uniform weights on the candidate support.
type reference_weights:: Optional[Float[Array, 'M']], default: None
param n_mosaic_points:: Quadrature points for mosaic broadening. Default: 7
type n_mosaic_points:: Union[int, Integer[Array, '']], default: 7
returns:: loss – Scalar objective value.
rtype:: Float[Array, '']