Index

A

aberration() (in module ptyrodactyl.electrons.forward) abstract_mesh (ptyrodactyl.tools.parallel.Mesh property) adagrad_update() (in module ptyrodactyl.tools.optimizers) adam_update() (in module ptyrodactyl.tools.optimizers) addressable_devices (ptyrodactyl.tools.parallel.NamedSharding property)

Any (class in ptyrodactyl.tools.loss_functions) (class in ptyrodactyl.tools.optimizers) args (ptyrodactyl.electrons.forward.partial attribute) axis_names (ptyrodactyl.tools.parallel.Mesh attribute) axis_sizes (ptyrodactyl.tools.parallel.Mesh property) axis_types (ptyrodactyl.tools.parallel.Mesh attribute)

C

Callable (class in ptyrodactyl.electrons.inverse) (class in ptyrodactyl.tools.optimizers) cbed() (in module ptyrodactyl.electrons.forward) cbed_no_slice() (in module ptyrodactyl.electrons.forward) center_thickness (ptyrodactyl.optics.lenses.LensParams attribute) check_compatible_aval() (ptyrodactyl.tools.parallel.NamedSharding method) Complex (class in ptyrodactyl.electrons.classes) (class in ptyrodactyl.electrons.forward) (class in ptyrodactyl.electrons.inverse) (class in ptyrodactyl.optics.epie) (class in ptyrodactyl.optics.lenses) (class in ptyrodactyl.tools.loss_functions) (class in ptyrodactyl.tools.optimizers)

complex_adagrad() (in module ptyrodactyl.tools.optimizers) complex_adam() (in module ptyrodactyl.tools.optimizers) complex_rmsprop() (in module ptyrodactyl.tools.optimizers) count() (ptyrodactyl.tools.optimizers.Sequence method) create_cosine_scheduler() (in module ptyrodactyl.tools.optimizers) create_lens_phase() (in module ptyrodactyl.optics.lenses) create_loss_function() (in module ptyrodactyl.tools.loss_functions) create_step_scheduler() (in module ptyrodactyl.tools.optimizers) create_warmup_cosine_scheduler() (in module ptyrodactyl.tools.optimizers)

D

device_ids (ptyrodactyl.tools.parallel.Mesh property) device_set (ptyrodactyl.tools.parallel.NamedSharding property) devices (ptyrodactyl.tools.parallel.Mesh attribute) diameter (ptyrodactyl.optics.lenses.LensParams attribute) Dict (in module ptyrodactyl.electrons.inverse) double_concave_lens() (in module ptyrodactyl.optics.lenses) double_convex_lens() (in module ptyrodactyl.optics.lenses)

dtypes (ptyrodactyl.electrons.classes.Complex attribute) (ptyrodactyl.electrons.forward.Complex attribute) (ptyrodactyl.electrons.inverse.Complex attribute) (ptyrodactyl.optics.epie.Complex attribute) (ptyrodactyl.optics.lenses.Complex attribute) (ptyrodactyl.tools.loss_functions.Complex attribute) (ptyrodactyl.tools.optimizers.Complex attribute)

E

empty (ptyrodactyl.tools.parallel.Mesh property)

F

focal_length (ptyrodactyl.optics.lenses.LensParams attribute) fourier_calib() (in module ptyrodactyl.electrons.forward)

fourier_coords() (in module ptyrodactyl.electrons.forward) func (ptyrodactyl.electrons.forward.partial attribute)

G

get_optimizer() (in module ptyrodactyl.electrons.inverse)

I

index() (ptyrodactyl.tools.optimizers.Sequence method) init (ptyrodactyl.tools.optimizers.Optimizer attribute) init_adagrad() (in module ptyrodactyl.tools.optimizers) init_adam() (in module ptyrodactyl.tools.optimizers) init_rmsprop() (in module ptyrodactyl.tools.optimizers) init_scheduler_state() (in module ptyrodactyl.tools.optimizers)

initial_lr (ptyrodactyl.tools.optimizers.LRSchedulerState attribute) initialize_mixed_states() (in module ptyrodactyl.electrons.inverse) initialize_probe_modes() (in module ptyrodactyl.electrons.inverse) initialize_random_modes() (in module ptyrodactyl.electrons.forward) is_fully_addressable (ptyrodactyl.tools.parallel.NamedSharding property) is_fully_replicated (ptyrodactyl.tools.parallel.NamedSharding property) is_multi_process (ptyrodactyl.tools.parallel.Mesh property)

J

jaxtyped() (in module ptyrodactyl.electrons.forward) (in module ptyrodactyl.optics.epie) (in module ptyrodactyl.optics.lenses)

K

keywords (ptyrodactyl.electrons.forward.partial attribute)

L

learning_rate (ptyrodactyl.tools.optimizers.LRSchedulerState attribute) lens_focal_length() (in module ptyrodactyl.optics.lenses) lens_thickness_profile() (in module ptyrodactyl.optics.lenses)

LensParams (class in ptyrodactyl.optics.lenses) local_devices (ptyrodactyl.tools.parallel.Mesh property) local_mesh (ptyrodactyl.tools.parallel.Mesh property) LRSchedulerState (class in ptyrodactyl.tools.optimizers)

M

m (ptyrodactyl.tools.optimizers.OptimizerState attribute) make_probe() (in module ptyrodactyl.electrons.forward) memory_kind (ptyrodactyl.tools.parallel.NamedSharding property) meniscus_lens() (in module ptyrodactyl.optics.lenses) Mesh (class in ptyrodactyl.tools.parallel) mesh (ptyrodactyl.tools.parallel.NamedSharding property) MixedQuantumStates (class in ptyrodactyl.electrons.classes) MixedStateParams (class in ptyrodactyl.electrons.classes) mode_weights (ptyrodactyl.electrons.classes.MixedStateParams attribute) modes (ptyrodactyl.electrons.classes.ProbeState attribute)

module ptyrodactyl.electrons.classes ptyrodactyl.electrons.forward ptyrodactyl.electrons.inverse ptyrodactyl.optics.epie ptyrodactyl.optics.lenses ptyrodactyl.tools.loss_functions ptyrodactyl.tools.optimizers ptyrodactyl.tools.parallel multi_mode_ptychography() (in module ptyrodactyl.electrons.inverse) multi_slice_ptychography() (in module ptyrodactyl.electrons.inverse)

N

n (ptyrodactyl.optics.lenses.LensParams attribute) NamedSharding (class in ptyrodactyl.tools.parallel) NamedTuple() (in module ptyrodactyl.electrons.classes) (in module ptyrodactyl.electrons.forward) (in module ptyrodactyl.electrons.inverse) (in module ptyrodactyl.optics.epie) (in module ptyrodactyl.optics.lenses) (in module ptyrodactyl.tools.optimizers)

normalize_mode_weights() (in module ptyrodactyl.electrons.forward) num_devices (ptyrodactyl.tools.parallel.NamedSharding property) num_modes (ptyrodactyl.electrons.classes.MixedStateParams attribute)

O

Optimizer (class in ptyrodactyl.tools.optimizers)

OptimizerState (class in ptyrodactyl.tools.optimizers)

P

partial (class in ptyrodactyl.electrons.forward) PartitionSpec (class in ptyrodactyl.tools.parallel) plano_concave_lens() (in module ptyrodactyl.optics.lenses) plano_convex_lens() (in module ptyrodactyl.optics.lenses) probabilities (ptyrodactyl.electrons.classes.MixedQuantumStates attribute) ProbeState (class in ptyrodactyl.electrons.classes) propagate_through_lens() (in module ptyrodactyl.optics.lenses) propagation_func() (in module ptyrodactyl.electrons.forward) ptyrodactyl.electrons.classes module ptyrodactyl.electrons.forward module ptyrodactyl.electrons.inverse module

ptyrodactyl.optics.epie module ptyrodactyl.optics.lenses module ptyrodactyl.tools.loss_functions module ptyrodactyl.tools.optimizers module ptyrodactyl.tools.parallel module Python Enhancement Proposals PEP 484 PEP 563 PyTree (class in ptyrodactyl.tools.loss_functions)

R

R1 (ptyrodactyl.optics.lenses.LensParams attribute) R2 (ptyrodactyl.optics.lenses.LensParams attribute) register_pytree_node_class() (in module ptyrodactyl.electrons.classes) (in module ptyrodactyl.optics.epie) (in module ptyrodactyl.optics.lenses)

rmsprop_update() (in module ptyrodactyl.tools.optimizers)

S

Sequence (class in ptyrodactyl.tools.optimizers) shape (ptyrodactyl.tools.parallel.Mesh property) shape_tuple (ptyrodactyl.tools.parallel.Mesh property) shard_array() (in module ptyrodactyl.tools.parallel) shift_beam_fourier() (in module ptyrodactyl.electrons.forward) single_slice_poscorrected() (in module ptyrodactyl.electrons.inverse) single_slice_ptychography() (in module ptyrodactyl.electrons.inverse)

size (ptyrodactyl.tools.parallel.Mesh property) spec (ptyrodactyl.tools.parallel.NamedSharding property) states (ptyrodactyl.electrons.classes.MixedQuantumStates attribute) stem_4d() (in module ptyrodactyl.electrons.forward) stem_4d_mixed_state() (in module ptyrodactyl.electrons.forward) stem_4d_multi() (in module ptyrodactyl.electrons.forward) step (ptyrodactyl.tools.optimizers.LRSchedulerState attribute) (ptyrodactyl.tools.optimizers.OptimizerState attribute)

T

transmission_func() (in module ptyrodactyl.electrons.forward) tree_flatten() (ptyrodactyl.electrons.classes.MixedQuantumStates method) (ptyrodactyl.electrons.classes.MixedStateParams method) (ptyrodactyl.electrons.classes.ProbeState method) (ptyrodactyl.optics.lenses.LensParams method) tree_unflatten() (ptyrodactyl.electrons.classes.MixedQuantumStates class method) (ptyrodactyl.electrons.classes.MixedStateParams class method) (ptyrodactyl.electrons.classes.ProbeState class method) (ptyrodactyl.optics.lenses.LensParams class method)

Tuple (in module ptyrodactyl.electrons.forward) (in module ptyrodactyl.electrons.inverse) (in module ptyrodactyl.optics.epie) (in module ptyrodactyl.optics.lenses) (in module ptyrodactyl.tools.optimizers) typechecker() (in module ptyrodactyl.electrons.forward)

U

UNCONSTRAINED (ptyrodactyl.tools.parallel.PartitionSpec attribute)

update (ptyrodactyl.tools.optimizers.Optimizer attribute)

V

v (ptyrodactyl.tools.optimizers.OptimizerState attribute)

W

wavelength_ang() (in module ptyrodactyl.electrons.forward) weights (ptyrodactyl.electrons.classes.ProbeState attribute)

wirtinger_grad() (in module ptyrodactyl.tools.optimizers) with_memory_kind() (ptyrodactyl.tools.parallel.NamedSharding method) with_spec() (ptyrodactyl.tools.parallel.NamedSharding method)