pyextal.jaxEngine.setParams¶

Classes¶

OptConfig

Configuration for optimization parameters

Functions¶

`prepare_reflection_bounds`() → Tuple[Tuple[Tuple[int, ...)	Prepare reflection bounds for optimization.
`denorm_fn`(p, lower, upper, m)
`norm_fn`(p, lower, upper, m)
`denormalize_params`(→ Dict)
`normalize_params`(→ Dict)
`get_bloch_constants`(→ Dict)	Extract constants from blochwave module with float32 precision
`calculate_tilt_dependent_values`(→ Dict)	Calculate all tilt-dependent values that don't change during optimization.
`calculate_static_ugh`(→ Tuple[jax.numpy.ndarray, ...)	Calculate Ugh contribution from static reflections.
`_check_thickness_range`(→ None)	Warns when `thickness_range` looks like Angstrom rather than a ratio.
`setup_optimization_params`(→ Dict)	Setup initial parameters and metadata for optimization.
`setup_optimization_model`(→ Dict)	Setup initial parameters and metadata for optimization.
`setup_structure_factors`(→ Tuple[Dict, List[Tuple[int, ...)
`fridel_symmetry_only`(→ Tuple[List[Tuple[int, int, ...)	Generate parameters dict and metadata for Friedel symmetry only case.
`spgsym`(hkl_all, sym[, sf])

Module Contents¶

class pyextal.jaxEngine.setParams.OptConfig¶

Bases: NamedTuple

Configuration for optimization parameters

One can choose to 1. normalize everything to a 0-1 range (set *_normalize=True) and optimize in the normalized space, then denormalize for the forward model. 2. normalize all other parameters to a range similar to structure factors but not the structure factors, (set beams_normalize=False, *_normalize= SOME SCALING FACTOR) There is no obvious difference in the result but the second approach is easier to extract results and no need to set up bounds for structure factors. Different learning rate might be needed for different approaches.

Args:

reflections: Tuple[Tuple[int, int, int], ...] | int¶

reflections_normalize: bool = False¶

reflections_upperbound: Tuple | None = None¶

reflections_lowerbound: Tuple | None = None¶

sf_mask: Tuple[bool, bool, bool, bool] = (True, True, True, True)¶

symmetry_constraint: str = 'Friedel'¶

loss_function: str = 'chi2'¶

batch_mode: str = 'chunking'¶

update_mode: str = 'per_epoch'¶

thickness_optimize: bool = False¶

initial_thickness: float = 100.0¶

thickness_range: Tuple[float, float] = (0.0, 200.0)¶

thickness_normalize: bool = True¶

probe_optimize: bool = False¶

initial_sigma_probe: float = 0.0¶

initial_gamma_probe: float = 0.0¶

probe_sigma_range: Tuple[float, float] = (0.0, 10.0)¶

probe_gamma_range: Tuple[float, float] = (0.0, 10.0)¶

probe_normalize: bool = True¶

use_voigt: bool = False¶

output_indices: Tuple[int, ...] = ()¶

output_shape: Tuple[int, int] = (0, 0)¶

chunk_shape: Tuple[int, int] = (0, 0)¶

kernel_radius: int = 0¶

stride: Tuple[int, int] = (0, 0)¶

learning_rate: float = 0.001¶

num_steps: int = 1000¶

warmup_steps: int = 0¶

fix_after: int = 0¶

initial_scale: float = 1.0¶

scale_range: Tuple[float, float] = (0.0, 10.0)¶

intensity_scale_optimize: bool = True¶

intensity_scale_normalize: bool = True¶

initial_background: float = 0.0¶

background_range: Tuple[float, float] = (0.0, 1.0)¶

background_optimize: bool = True¶

background_normalize: bool = True¶

power_regularization_weight: float = 0.0¶

nlayers: int = 1¶

frozen_layers: Tuple[int, ...] = ()¶

diagonalise_measured: bool = True¶

model_name: str | None = None¶

model_params: list[dict] | None = None¶

model_params_normalize: bool = True¶

model_params_range: list[dict] | None = None¶

refined_entity: str = 'structure_factor'¶

oversub_indices: Tuple[int, ...] = ()¶

oversub_leak_index: int = -1¶

oversub_optimize: bool = False¶

initial_oversub_offset: Tuple[float, ...] = ()¶

initial_oversub_leak: Tuple[float, ...] = ()¶

oversub_offset_upper: Tuple[float, ...] = ()¶

oversub_leak_upper: Tuple[float, ...] = ()¶

classmethod from_yaml(filepath: str) → OptConfig¶: Load OptConfig from YAML file

to_yaml(filepath: str) → None¶: Save OptConfig to YAML file in a human-readable format.

pyextal.jaxEngine.setParams.prepare_reflection_bounds(reflections: List[Tuple[int, int, int]], values: List[Tuple[float, float, float, float]] | None = None, upperbounds: List[Tuple[float, float, float, float]] | None = None, lowerbounds: List[Tuple[float, float, float, float]] | None = None, bound_percentage: tuple = (0.05, 0.05, 0.05, 0.05)) → Tuple[Tuple[Tuple[int, int, int], ...], Tuple[Tuple[float, float, float, float], ...], Tuple[Tuple[float, float, float, float], ...], Tuple[Tuple[float, float, float, float], ...]]¶

Prepare reflection bounds for optimization.

Parameters:

reflections – List of reflection tuples (h, k, l)
values – Optional list of initial values [U_amp, U_phase, UA_amp, UA_phase] for each reflection
upperbounds – Optional list of upper bounds for each reflection
lowerbounds – Optional list of lower bounds for each reflection
bound_percentage – Percentage to use for automatic bounds (default 5%)

Returns:

Tuple of (reflections, values, upperbounds, lowerbounds) as tuples for OptConfig

pyextal.jaxEngine.setParams.denorm_fn(p, lower, upper, m)¶

pyextal.jaxEngine.setParams.norm_fn(p, lower, upper, m)¶

pyextal.jaxEngine.setParams.denormalize_params(params: Dict, metadata: Dict, mask: Dict) → Dict¶

pyextal.jaxEngine.setParams.normalize_params(params: Dict, metadata: Dict, mask: Dict) → Dict¶

pyextal.jaxEngine.setParams.get_bloch_constants(alpha: float) → Dict¶: Extract constants from blochwave module with float32 precision

pyextal.jaxEngine.setParams.calculate_tilt_dependent_values(tilts: jax.numpy.ndarray, constants: Dict, config: OptConfig) → Dict¶

Calculate all tilt-dependent values that don’t change during optimization. This includes excitation errors (sg), bigkz, and bigkn values.

Parameters:

tilts – Array of tilt vectors, shape (n_tilts, 3)
constants – Dictionary containing hkl, gmxr, bigk, zone, alpha, snorm

Returns:

Dict containing pre-calculated values for all tilts

pyextal.jaxEngine.setParams.calculate_static_ugh(params_dict: Dict, reflection_list: List, phase: jax.numpy.ndarray, mask: jax.numpy.ndarray, nbeams: int) → Tuple[jax.numpy.ndarray, jax.numpy.ndarray]¶: Calculate Ugh contribution from static reflections.

pyextal.jaxEngine.setParams._check_thickness_range(config: OptConfig) → None¶

Warns when thickness_range looks like Angstrom rather than a ratio.

The forward model uses the thickness parameter as a ratio of initial_thickness, so the bounds are ratios too. Bounds in Angstrom make every normalised Adam step initial_thickness times too large.

pyextal.jaxEngine.setParams.setup_optimization_params(config: OptConfig, tilts: jax.numpy.ndarray, constants: Dict, initial_value: List[Dict] | None = None) → Dict¶: Setup initial parameters and metadata for optimization. :returns: Dict with ‘params’ (dict of dicts) and ‘metadata’ for reconstruction

pyextal.jaxEngine.setParams.setup_optimization_model(config: OptConfig, tilts: jax.numpy.ndarray, constants: Dict) → Dict¶: Setup initial parameters and metadata for optimization. :returns: Dict with ‘params’ (dict of dicts) and ‘metadata’ for reconstruction

pyextal.jaxEngine.setParams.setup_structure_factors(hkl_all: numpy.ndarray, sym: pyextal.symmetry.GroupSymmetry, config: OptConfig, sf: numpy.ndarray = None) → Tuple[Dict, List[Tuple[int, int, int]], jax.numpy.ndarray, jax.numpy.ndarray]¶

pyextal.jaxEngine.setParams.fridel_symmetry_only(hkl_all: numpy.ndarray, sf: numpy.ndarray = None) → Tuple[List[Tuple[int, int, int]], Dict]¶: Generate parameters dict and metadata for Friedel symmetry only case.

pyextal.jaxEngine.setParams.spgsym(hkl_all: numpy.ndarray, sym: pyextal.symmetry.GroupSymmetry, sf: numpy.ndarray = None)¶