def test_consistent_lstsq_grad(self):
"""Check ptycho.solver.lstsq_grad for consistency."""
algorithm = f"{'mpi-' if _mpi_size > 1 else ''}lstsq_grad-position"
result = self.template_consistent_algorithm(params={
'algorithm_options':
tike.ptycho.LstsqOptions(
num_batch=5,
num_iter=16,
),
'num_gpu':
2,
'position_options':
PositionOptions(
self.scan,
use_adaptive_moment=True,
),
'probe_options':
ProbeOptions(),
'object_options':
ObjectOptions(),
'use_mpi':
_mpi_size > 1,
}, )
_save_ptycho_result(result, algorithm)
self._save_position_error_variance(result, algorithm)
def test_consistent_rpie(self):
"""Check ptycho.solver.rpie for consistency."""
_save_ptycho_result(
self.template_consistent_algorithm(params={
'algorithm_options':
tike.ptycho.RpieOptions(
num_batch=5,
num_iter=16,
),
'num_gpu':
2,
'probe_options':
ProbeOptions(),
'object_options':
ObjectOptions(),
'use_mpi':
_mpi_size > 1,
}, ), f"{'mpi-' if _mpi_size > 1 else ''}rpie")
def test_consistent_lstsq_grad(self):
"""Check ptycho.solver.lstsq_grad for consistency."""
_save_ptycho_result(
self.template_consistent_algorithm(params={
'algorithm_options':
tike.ptycho.LstsqOptions(
num_batch=5,
num_iter=16,
),
'num_gpu':
2,
'probe_options':
ProbeOptions(orthogonality_constraint=True, ),
'object_options':
ObjectOptions(use_adaptive_moment=True, ),
'use_mpi':
_mpi_size > 1,
}, ), f"{'mpi-' if _mpi_size > 1 else ''}lstsq_grad")
def test_consistent_rpie_ref(self):
"""Check ptycho.solver.rpie position correction."""
algorithm = f"{'mpi-' if _mpi_size > 1 else ''}rpie-position-ref"
result = self.template_consistent_algorithm(params={
'algorithm_options':
tike.ptycho.RpieOptions(
num_batch=5,
num_iter=16,
),
'num_gpu':
2,
'probe_options':
ProbeOptions(),
'object_options':
ObjectOptions(),
'use_mpi':
_mpi_size > 1,
}, )
_save_ptycho_result(result, algorithm)
def test_consistent_lstsq_grad_variable_probe(self):
"""Check ptycho.solver.lstsq_grad for consistency."""
probes_with_modes = min(3, self.probe.shape[-3])
eigen_probe, weights = tike.ptycho.probe.init_varying_probe(
self.scan,
self.probe,
num_eigen_probes=3,
probes_with_modes=probes_with_modes,
)
result = self.template_consistent_algorithm(params={
'algorithm_options':
tike.ptycho.LstsqOptions(
num_batch=5,
num_iter=16,
),
'num_gpu':
2,
'probe_options':
ProbeOptions(),
'object_options':
ObjectOptions(use_adaptive_moment=True, ),
'use_mpi':
_mpi_size > 1,
'eigen_probe':
eigen_probe,
'eigen_weights':
weights,
}, )
_save_ptycho_result(
result,
f"{'mpi-' if _mpi_size > 1 else ''}lstsq_grad-variable-probe",
)
assert np.all(
result['eigen_weights'][..., 1:, probes_with_modes:] == 0), (
"These weights should be unused/untouched "
"and should have been initialized to zero.")