def _calculatePhaseModulationsForRCAngles(self): logger.info("Calculating the phase modulations for the rc angles.") ttheta = self.scan_grid.two_theta domega = self.scan_grid.del_omega ki = 2 * np.pi / self.wavelength * np.array( [np.cos(ttheta), np.sin(ttheta), 0]) kf = 2 * np.pi / self.wavelength * np.array([1, 0, 0]) q = (kf - ki)[:, None] ki_new = 2 * np.pi / self.wavelength * np.array([ np.cos(ttheta + self.scan_grid.rc_angles), np.sin(ttheta + self.scan_grid.rc_angles), 0 * self.scan_grid.rc_angles ]) kf_new = 2 * np.pi / self.wavelength * np.array([ np.cos(self.scan_grid.rc_angles), np.sin(self.scan_grid.rc_angles), 0 * self.scan_grid.rc_angles ]) q_new = kf_new - ki_new delta_q = q_new - q # Probe dimensions in real space (assumes even shape) position_grids = [ np.arange(-s // 2, s // 2) * ds for (s, ds) in zip(self.probe_3d.shape, self.probe_3d.pixel_size) ] Ry, Rx, Rz = np.meshgrid(*position_grids, indexing='ij') phase_modulations_all = np.exp(1j * np.array([ delta_q[0, i] * Ry + delta_q[1, i] * Rx + delta_q[2, i] * Rz for i in range(self.scan_grid.n_rc_angles) ])) self._phase_modulations_all = phase_modulations_all
def __init__(self, obj: Obj, probe: Probe3D, scan_grid: BraggPtychoGrid, exit_wave_axis: str = "y", upsampling_factor: int = 1, setup_second_order: bool = False, dtype: str = "float32"): if scan_grid.full_field_probe: e = ValueError( "Full field probe not supported for Bragg ptychography.") logger.error(e) raise e if scan_grid.grid2d_axes != ("y", "z"): e = ValueError( "Only supports the case where the ptychographic scan is on the yz-plane." ) logger.error(e) raise e if exit_wave_axis != 'y': e = ValueError( "Only supports the case where the exit waves are output along the y-direction." ) logger.error(e) raise e super().__init__(obj, probe, scan_grid, upsampling_factor, setup_second_order, dtype) logger.info("Creating the phase modulations for the scan angles.") with tf.device("/gpu:0"): self._probe_phase_modulations_all_t = tf.constant( self._getProbePhaseModulationsStack(), dtype='complex64') self._full_rc_positions_indices_t = tf.constant( scan_grid.full_rc_positions_indices, dtype='int64')
def createObj(obj_params_dict: dict, det_3d_shape: Tuple[int, int, int], obj_pixel_size: Tuple[float, float, float], upsampling_factor: float = 1.0) -> Obj: logger.info("Creating new crystal cell.") while True: obj_crystal = Simulated3DCrystalCell(**obj_params_dict) if det_3d_shape[ 1] * upsampling_factor >= obj_crystal.bordered_array.shape[ 1]: break else: logger.warning( "Generated object width is larger than detector x-width. Trying again." ) logger.info( 'Adding x and z borders to crystal cell based on detector parameters.' ) # Ensuring that the number of pixels in the x dimension matches that in the detector. padx = (det_3d_shape[1] - obj_crystal.bordered_array.shape[1]) // 2 # Calculating the padding needed to accomodate all feasible probe translations (w probe-obj interaction) padz, _ = Simulation.calculateObjBorderAfterRotation( 0, det_3d_shape[2], obj_crystal.bordered_array.shape[1], 0) pad_shape = (0, padx, padz) obj = CustomObjFromArray(array=obj_crystal.bordered_array, border_shape=np.vstack( (pad_shape, pad_shape)).T, border_const=0, pixel_size=obj_pixel_size) return obj
def getFlopsPerIter(self) -> int: if self.n_validation > 0: e = RuntimeWarning("n_validation > 0 can give misleading flops.") logger.warning(e) if hasattr(self, "session"): e = RuntimeWarning( "Calculating computational cost after previously initializing the static graph can " + "include inessential calculations (like training and/or validation loss value) " + "and thus give misleading results.") logger.warning(e) else: with self.graph.as_default(): config = tf.ConfigProto() config.gpu_options.allow_growth = True logger.info("Initializing the session.") self.session = tf.Session(config=config) self.session.run(tf.global_variables_initializer()) for opt in self.optimizers: self.session.run(opt.minimize_op) total_flops = getComputationalCostInFlops(self.graph) return total_flops
def finalize(self): columns = [] for item in self._datalog_items: columns.append(item.title) logger.info("Initializing the log outputs...") self.dataframe = DataFrame(columns=columns, dtype='float32') self.dataframe.loc[0] = np.nan
def _initGraph(self): self.graph = tf.Graph() with self.graph.as_default(): with tf.device("/gpu:0"): self._amplitudes_t = tf.constant(self.amplitudes, dtype=self.dtype) if self.intensities_mask is not None: self._data_mask_t = tf.constant(self.intensities_mask, dtype=tf.bool) logger.info('creating batches...') self._createDataBatches()
def createProbe3D(probe_params_dict: dict, wavelength: float, rotate_angle: float, probe_xz_shape: Tuple[int, int], obj_pixel_size: Tuple[int, int, int]) -> Probe: logger.info("Creating new guassian, speckled, 2d probe.") probe_yz = GaussianSpeckledProbe(wavelength=wavelength, **probe_params_dict) ny = probe_yz.shape[0] nx, nz = probe_xz_shape # overdoing the repeat and interpolation just for safety logger.info( "Rotating and interpolating the 2d probe to generate the 3d probe." ) probe_yz_stack = np.repeat(probe_yz.wavefront.ifftshift[:, None, :], nx * 2, axis=1) rdeg = rotate_angle * 180 / np.pi rotated_real = ndimage.rotate(np.real(probe_yz_stack), rdeg, axes=(0, 1), mode='constant', order=1) rotated_imag = ndimage.rotate(np.imag(probe_yz_stack), rdeg, axes=(0, 1), mode='constant', order=1) rotated = rotated_real + 1j * rotated_imag # Calculating the number of pixels required to capture the y-structure of the rotated probe rny = ny / np.cos(rotate_angle) // 1 rshape = np.array([rny, nx, nz]).astype('int') # Calculating the extent of the probe (relative to the center of rotation) in the x and y dimensions and # also ensuring that the probe array has an even number of pixels. # Adding any required padding so that the z dimension of the probe matches up with the detector. a = (np.array(rotated.shape) - rshape) // 2 b = ((np.array(rotated.shape) - rshape) // 2 + (np.array(rotated.shape) - rshape) % 2) z_pad = (rshape[2] - rotated.shape[2]) // 2 rotated_centered = rotated[a[0]:-b[0], a[1]:-b[1]] rotated_centered = np.pad(rotated_centered, [[0, 0], [0, 0], [z_pad, z_pad]], mode='constant', constant_values=0) probe_3d = CustomProbe3DFromArray(array=rotated_centered, wavelength=wavelength, pixel_size=obj_pixel_size) return probe_3d
def __init__(self, *args: int, obj_array_true: np.ndarray = None, probe_wavefront_true: np.ndarray = None, **kwargs: int): logger.info('initializing...') super().__init__(*args, **kwargs) logger.info('attaching fwd model...') self.attachForwardModel("farfield") logger.info('creating loss fn...') self.attachLossFunction("least_squared") logger.info('creating optimizers...') self.attachOptimizerPerVariable( "obj", optimizer_type="adam", optimizer_init_args={"learning_rate": 1e-2}) self.attachOptimizerPerVariable( "probe", optimizer_type="adam", optimizer_init_args={"learning_rate": 1e-1}, initial_update_delay=0) if obj_array_true is not None: self.addCustomMetricToDataLog( title="obj_error", tensor=self.fwd_model.obj_cmplx_t, log_epoch_frequency=10, registration_ground_truth=obj_array_true) if probe_wavefront_true is not None: self.addCustomMetricToDataLog( title="probe_error", tensor=self.fwd_model.probe_cmplx_t, log_epoch_frequency=10, registration_ground_truth=probe_wavefront_true)
def __init__(self, *args: int, loss_type: str = "least_squared", obj_array_true: np.ndarray = None, learning_rate_obj: float = 1e-2, registration_log_frequency=10, **kwargs): logger.info('initializing...') super().__init__(*args, **kwargs) logger.info('attaching fwd model...') self.attachForwardModel("bragg") logger.info('creating loss fn...') self.attachLossFunction(loss_type) logger.info('creating optimizers...') self.attachTensorflowOptimizerForVariable( "obj", optimizer_type="adam", optimizer_init_args={"learning_rate": learning_rate_obj}) if obj_array_true is not None: self.addCustomMetricToDataLog( title="obj_error", tensor=self.fwd_model.obj_cmplx_t, log_epoch_frequency=registration_log_frequency, registration_ground_truth=obj_array_true)
def getFlopsPerIter(self) -> dict: if self.n_validation > 0: e = RuntimeWarning("n_validation > 0 can give misleading flops.") logger.warning(e) if hasattr(self, "session"): e = RuntimeWarning( "Calculating computational cost after previously initializing the static graph can " + "include inessential calculations (like training and/or validation loss value) " + "and thus give misleading results.") logger.warning(e) else: with self.graph.as_default(): config = tf.ConfigProto() config.gpu_options.allow_growth = True logger.info("Initializing the session.") self.session = tf.Session(config=config) self.session.run(tf.global_variables_initializer()) for opt in self.optimizers: self.session.run(opt.minimize_op) total_flops = getComputationalCostInFlops(self.graph) ls_flops = getComputationalCostInFlops(self.graph, keywords=[ ("opt_minimize_step", "backtracking_linesearch") ], exclude_keywords=False) opt_only_flops = getComputationalCostInFlops(self.graph, keywords=["opt"], exclude_keywords=False) flops_without_ls = total_flops - ls_flops d = { "total_flops": total_flops, "obj_ls_flops": ls_flops, "probe_ls_flops": 0, "obj_only_flops": opt_only_flops, "probe_only_flops": 0, "flops_without_ls": flops_without_ls } return d
def __init__(self, obj: ptychoSampling.obj.Obj, probe: ptychoSampling.probe.Probe, grid: ptychoSampling.grid.ScanGrid, intensities: np.ndarray, intensities_mask: np.ndarray = None, n_validation: int = 0, training_batch_size: int = 0, validation_batch_size: int = 0, background_level: float = 0., dtype: str = 'float32'): self.obj = copy.deepcopy(obj) self.probe = copy.deepcopy(probe) self.grid = copy.deepcopy(grid) self.amplitudes = intensities**0.5 self.background_level = background_level if intensities_mask is not None: if intensities_mask.dtype != np.bool: raise ValueError("Mask supplied should be a boolean array.") self.intensities_mask = intensities_mask self.dtype = dtype if self.dtype == "float32": self._eps = 1e-8 elif self.dtype == "float64": self._eps = 1e-8 else: raise ValueError self._splitTrainingValidationData(n_validation, training_batch_size, validation_batch_size) self._initGraph() logger.info('creating pandas log...') self.iteration = 0 self.datalog = DataLogs() self._default_log_items = {"epoch": None, "train_loss": None} for key in self._default_log_items: self.datalog.addSimpleMetric(key) if self.n_validation > 0: self._validation_log_items = { "validation_loss": None, "validation_min": None, "patience": None } for key in self._validation_log_items: self.datalog.addSimpleMetric(key)
def createScanGrid(scan_grid_params_dict: dict, angle_grid_params_dict: dict, bordered_obj_yz_shape: Tuple[int, int], probe_yz_shape: Tuple[int, int]) -> BraggPtychoGrid: logger.info( "creating new 2d scan grid based on object and probe shapes.") scan_grid_2d = RectangleGrid(obj_w_border_shape=bordered_obj_yz_shape, probe_shape=probe_yz_shape, **scan_grid_params_dict) scan_grid_2d.checkOverlap() logger.info("Using created 2d scan grid to create full RC scan grid.") scan_grid = BraggPtychoGrid.fromPtychoScan2D(scan_grid_2d, grid2d_axes=("y", "z"), **angle_grid_params_dict) return scan_grid
def __init__(self, obj: Obj, probe: Probe, scan_grid: ScanGrid, upsampling_factor: int = 1): self.model_vars = {} with tf.device("/gpu:0"): self.obj_cmplx_t = self._addComplexVariable(obj.array, name="obj") self.obj_w_border_t = tf.pad(self.obj_cmplx_t, obj.border_shape, constant_values=obj.border_const) self.probe_cmplx_t = self._addComplexVariable(probe.wavefront, "probe") self.upsampling_factor = upsampling_factor logger.info("Creating obj views for the scan positions.") self._obj_views_all_t = self._getPtychoObjViewStack(obj, probe, scan_grid)
def _calculateGaussianFit(self) -> None: r"""Fit a 2d gaussian to the probe intensities. Performs a least-squares fit (using ``scipy.optimize.curve_fit``) to fit a 2d gaussian to the probe intensities. Uses the calculated gaussian spread to calculate the FWHM as well. See also -------- gaussian_fit utils.generalized2dGaussian """ logger.info( 'Fitting a generalized 2d gaussian to the probe intensity.') from scipy.optimize import curve_fit nx = self.shape[-1] ny = self.shape[-2] #intensities = np.fft.ifftshift(np.abs(self.wavefront)**2) intensities = self.wavefront.fftshift.intensities y = np.arange(-ny // 2, ny // 2) * self.pixel_size[0] x = np.arange(-ny // 2, nx // 2) * self.pixel_size[1] yy, xx = np.meshgrid(y, x) xdata = np.stack((xx.flatten(), yy.flatten()), axis=1) bounds_min = [0, x[0], y[0], 0, 0, -np.pi / 4, 0] bounds_max = [ intensities.sum(), x[-1], y[-1], x[-1] * 2, y[-1] * 2, np.pi / 4, intensities.max() ] popt, _ = curve_fit(utils.generalized2dGaussian, xdata, intensities.flatten(), bounds=[bounds_min, bounds_max]) amplitude, center_x, center_y, sigma_x, sigma_y, theta, offset = popt self._gaussian_fit = { "amplitude": amplitude, "center_x": center_x, "center_y": center_y, "sigma_x": sigma_x, "sigma_y": sigma_y, "theta": theta, "offset": offset } self._gaussian_fwhm = 2.355 * np.array((sigma_y, sigma_x))
def _getPtychoObjViewStack( self, obj_w_border_t: tf.Tensor, probe_cmplx_t: tf.Tensor, position_indices_t: tf.Tensor = None) -> tf.Tensor: """Precalculate the object positioning for each scan position. Assumes a small object that is translated within the dimensions of a full-field probe. For each scan position, we translate the object, then pad the object array to the size of the probe beam. For the padding, we assume free-space (transparent) propagation and use 1.0. In Tensorflow, performing the pad-and-stack procedure in the GPU for complex -valued arrays seems to be buggy. As a workaround, we separately pad-and-stack the real and imaginary parts of the object with 1.0 and 0 respectively. Returns ---------- obj_views : tensor(complex) Stack of tensors that correspond to the padded object at each object translation. """ if position_indices_t is None: position_indices_t = tf.range( self._scan_grid.positions_pix.shape[0]) logger.info("Creating obj views for the scan positions.") if not hasattr(self, "_obj_view_indices_t"): obj_view_indices = self._genViewIndices( self._scan_grid.positions_pix) self._obj_view_indices_t = tf.constant(obj_view_indices) batch_obj_view_indices_t = tf.gather(self._obj_view_indices_t, position_indices_t) batch_obj_views_t = tf.gather(tf.reshape(obj_w_border_t, [-1]), batch_obj_view_indices_t) #obj_view_indices = self._genViewIndices(scan_grid.positions_pix) #obj_view_indices_t = tf.constant(obj_view_indices, dtype='int64') #obj_views_t = tf.gather(tf.reshape(self.obj_w_border_t, [-1]), obj_view_indices_t) #obj_views_t = tf.reshape(obj_views_t, # (obj_view_indices.shape[0], # *(self.probe_cmplx_t.get_shape().as_list()))) return batch_obj_views_t
def __init__(self, obj: ptychoSampling.obj.Obj, probe: ptychoSampling.probe.Probe, grid: ptychoSampling.grid.ScanGrid, intensities: np.ndarray, n_validation: int = 0, batch_size: int = 0): self.obj = copy.deepcopy(obj) self.probe = copy.deepcopy(probe) self.grid = copy.deepcopy(grid) self.amplitudes = intensities**0.5 self.n_all = self.amplitudes.shape[0] self.n_validation = n_validation self.n_train = self.n_all - self.n_validation self.batch_size = batch_size self.graph = tf.Graph() with self.graph.as_default(): with tf.device("/gpu:0"): self._amplitudes_t = tf.constant(self.amplitudes, dtype=tf.float32) logger.info('creating batches...') self._createDataBatches() logger.info('creating log...') self.iteration = 0 self.datalog = DataLogs() self._default_log_items = {"epoch": None, "train_loss": None} for key in self._default_log_items: self.datalog.addSimpleMetric(key) if self.n_validation > 0: self._validation_log_items = { "validation_loss": None, "validation_min": None, "patience": None } for key in self._validation_log_items: self.datalog.addSimpleMetric(key)
def _calculateDiffractionPatterns(self): # Calculating the wave vectors self._calculatePhaseModulationsForRCAngles() intensities_all = [] logger.info("Calculating the generated diffraction patterns.") for ia in range(self.scan_grid.n_rc_angles): for ib, (py, pz) in enumerate(self.scan_grid.positions_pix): obj_slice = self.obj.bordered_array[py:py + self.probe_3d.shape[0], :, pz:pz + self.probe_3d.shape[2]] exit_wave = obj_slice * self.probe_3d.wavefront * self._phase_modulations_all[ ia] exit_wave_proj = np.sum(exit_wave, axis=0).fftshift intensities_all.append(exit_wave_proj.propFF().intensities) self.intensities = np.random.poisson( intensities_all) if self.poisson_noise else np.array( intensities_all)
def finalizeSetup(self): self._checkAttr("optimizers", "finalize") logger.info("finalizing the data logger.") self.datalog.finalize() with self.graph.as_default(): config = tf.ConfigProto() config.gpu_options.allow_growth = True logger.info("Initializing the session.") self.session = tf.Session(config=config) self.session.run(tf.global_variables_initializer()) logger.info("Finalized setup.")
def __init__(self, *args: int, obj_array_true: np.ndarray = None, probe_wavefront_true: np.ndarray = None, shuffle_order: list = None, **kwargs: int): logger.info('initializing...') super().__init__(*args, **kwargs) logger.info('attaching fwd model...') self.attachForwardModel("farfield") logger.info('creating loss fn...') self.attachLossFunction("least_squared")
def __init__(self, *args: int, loss_type: str = "least_squared", obj_array_true: np.ndarray = None, probe_wavefront_true: np.ndarray = None, update_delay_probe: int = 0, update_delay_obj: int = 0, reconstruct_probe: bool = True, registration_log_frequency: int = 10, opt_init_extra_kwargs: dict = None, obj_abs_proj: bool = True, loss_init_extra_kwargs: dict = None, r_factor_log: bool = True, apply_precond: bool = False, **kwargs: int): logger.info('initializing...') super().__init__(*args, **kwargs) if self.training_batch_size != self.n_train: e = ValueError( "Conjugate gradient reconstruction does not support minibatch reconstruction." ) logger.error(e) raise e logger.info('attaching fwd model...') self._attachCustomForwardModel(JointFarfieldForwardModelT, obj_abs_proj=obj_abs_proj) logger.info('creating loss fn...') self.attachLossFunctionSecondOrder( loss_type, loss_init_extra_kwargs=loss_init_extra_kwargs) logger.info("creating optimizers...") if opt_init_extra_kwargs is None: opt_init_extra_kwargs = {} self._preconditioner = None if apply_precond: with self.graph.as_default(): loss_data_type = self._loss_method.data_type hessian_t = ( tf.ones_like(self._batch_train_predictions_t) * self._train_hessian_fn(self._batch_train_predictions_t)) hessian_t = tf.reshape(hessian_t, [-1, *self.probe.shape]) obj_scaling = self._getObjScaling(loss_data_type, hessian_t) probe_scaling = self._getProbeScaling(loss_data_type, hessian_t) scaling_both = tf.concat((obj_scaling, probe_scaling), axis=0) scaling_both = tf.concat((scaling_both, scaling_both), axis=0) zero_condition = tf.less(scaling_both, 1e-10 * tf.reduce_max(scaling_both)) zero_case = tf.ones_like(scaling_both) / ( 1e-8 * tf.reduce_max(scaling_both)) self._preconditioner = tf.where(zero_condition, zero_case, 1 / scaling_both) opt_init_args = { "input_var": self.fwd_model.joint_v, "predictions_fn": self._predictions_fn, "loss_fn": self._train_loss_fn, "name": "opt", "diag_precondition_t": self._preconditioner } opt_init_args.update(opt_init_extra_kwargs) self._attachCustomOptimizerForVariable( ConjugateGradientOptimizer, optimizer_init_args=opt_init_args) self.addCustomMetricToDataLog( title="ls_iters", tensor=self.optimizers[0]._optimizer._linesearch_steps, log_epoch_frequency=1) self.addCustomMetricToDataLog( title="alpha", tensor=self.optimizers[0]._optimizer._linesearch._alpha, log_epoch_frequency=1) if obj_array_true is not None: self.addCustomMetricToDataLog( title="obj_error", tensor=self.fwd_model.obj_cmplx_t, log_epoch_frequency=registration_log_frequency, registration_ground_truth=obj_array_true) if reconstruct_probe and (probe_wavefront_true is not None): self.addCustomMetricToDataLog( title="probe_error", tensor=self.fwd_model.probe_cmplx_t, log_epoch_frequency=registration_log_frequency, registration_ground_truth=probe_wavefront_true) self._addRFactorLog(r_factor_log, registration_log_frequency)
def __init__( self, *args: int, loss_type: str = "least_squared", obj_array_true: np.ndarray = None, probe_wavefront_true: np.ndarray = None, update_delay_probe: int = 0, update_frequency_probe: int = 1, update_frequency_obj: int = 1, reconstruct_probe: bool = True, registration_log_frequency: int = 1, opt_init_extra_kwargs: dict = None, obj_abs_proj: bool = True, loss_init_extra_kwargs: dict = None, r_factor_log: bool = True, #### These two parameters are experimental only. Not to be used in production simulations. # apply_diag_mu_scaling: bool = False, # Does not help apply_precond: bool = False, ########################################################################################### **kwargs: int): print("opt_init_extra_kwargs", opt_init_extra_kwargs) print("Loss init args", loss_init_extra_kwargs) print("update_delay_probe", update_delay_probe, "update_frequency", update_frequency_probe) print("update_frequency_obj", update_frequency_obj) logger.info('initializing...') super().__init__(*args, **kwargs) logger.info('attaching fwd model...') self._attachCustomForwardModel(JointFarfieldForwardModelT, obj_abs_proj=obj_abs_proj) logger.info('creating loss fn...') self.attachLossFunctionSecondOrder( loss_type, loss_init_extra_kwargs=loss_init_extra_kwargs) logger.info("creating optimizers...") if opt_init_extra_kwargs is None: opt_init_extra_kwargs = {} opt_init_args = { "input_var": self.fwd_model.joint_v, "predictions_fn": self._predictions_fn, "loss_fn": self._train_loss_fn, "diag_hessian_fn": self._train_hessian_fn, "damping_factor": 10.0, "damping_update_frequency": 5, "damping_update_factor": 0.99, "name": "opt" } # Experimental only.######################################################################################## if apply_precond: loss_data_type = self._loss_method.data_type with self.graph.as_default(): hessian_t = ( tf.ones_like(self._batch_train_predictions_t) * self._train_hessian_fn(self._batch_train_predictions_t)) hessian_t = tf.reshape(hessian_t, [-1, *self.probe.shape]) obj_scaling_t = self._getObjScaling(loss_data_type, hessian_t) probe_scaling_t = self._getProbeScaling( loss_data_type, hessian_t) scaling_both = tf.concat((obj_scaling_t, probe_scaling_t), axis=0) self._joint_scaling_t = tf.concat((scaling_both, scaling_both), axis=0) opt_init_args['diag_precond_t'] = self._joint_scaling_t ############################################################################################################# opt_init_args.update(opt_init_extra_kwargs) self._attachCustomOptimizerForVariable( CurveballOptimizer, optimizer_init_args=opt_init_args) self.addCustomMetricToDataLog( title="mu", tensor=self.optimizers[0]._optimizer._damping_factor, log_epoch_frequency=1) if obj_array_true is not None: self.addCustomMetricToDataLog( title="obj_error", tensor=self.fwd_model.obj_cmplx_t, log_epoch_frequency=registration_log_frequency, registration_ground_truth=obj_array_true) if (probe_wavefront_true is not None): self.addCustomMetricToDataLog( title="probe_error", tensor=self.fwd_model.probe_cmplx_t, log_epoch_frequency=registration_log_frequency, registration_ground_truth=probe_wavefront_true) self._addRFactorLog(r_factor_log, registration_log_frequency)
def __init__(self, *args: int, loss_type: str = "least_squared", obj_array_true: np.ndarray = None, probe_wavefront_true: np.ndarray = None, max_cg_iter: int = 100, min_cg_tol: float = 0.1, registration_log_frequency: int = 10, both_registration_nlse: bool = True, opt_init_extra_kwargs: dict = None, obj_abs_proj: bool = True, obj_abs_max: float = 1.0, probe_abs_max: float = None, loss_init_extra_kwargs: dict = None, r_factor_log: bool = True, update_delay_probe: int = 0, reconstruct_probe: bool = True, apply_diag_mu_scaling: bool = True, apply_precond: bool = False, apply_precond_and_scaling: bool = False, stochastic_diag_estimator_type: str = None, stochastic_diag_estimator_iters: int = 1, **kwargs: int): logger.info('initializing...') super().__init__(*args, **kwargs) if apply_precond_and_scaling: logger.info( "Overriding any set values for apply_precond and apply_diag_mu_scaling" ) apply_precond = True apply_diag_mu_scaling = True if stochastic_diag_estimator_type is not None: if apply_diag_mu_scaling or apply_precond: e = ValueError( "Cannot use analytical precond and diag ggn elems if stochastic calculation is enabled." ) logger.error(e) raise e if self.training_batch_size != self.n_train: e = ValueError( "LMA reconstruction does not support minibatch reconstruction." ) logger.error(e) raise e logger.info('attaching fwd model...') self._attachCustomForwardModel(JointFarfieldForwardModelT, obj_abs_proj=obj_abs_proj, obj_abs_max=obj_abs_max, probe_abs_max=probe_abs_max) logger.info('creating loss fn...') print('Loss init args', loss_init_extra_kwargs) self.attachLossFunctionSecondOrder( loss_type, loss_init_extra_kwargs=loss_init_extra_kwargs) logger.info("creating optimizer...") if opt_init_extra_kwargs is None: opt_init_extra_kwargs = {} if apply_diag_mu_scaling or apply_precond: with self.graph.as_default(): loss_data_type = self._loss_method.data_type hessian_t = ( tf.ones_like(self._batch_train_predictions_t) * self._train_hessian_fn(self._batch_train_predictions_t)) hessian_t = tf.reshape(hessian_t, [-1, *self.probe.shape]) print(hessian_t) self._obj_scaling = self._getObjScaling( loss_data_type, hessian_t) self._probe_scaling = self._getProbeScaling( loss_data_type, hessian_t) scaling_both = tf.concat( (self._obj_scaling, self._probe_scaling), axis=0) if apply_diag_mu_scaling: self._joint_scaling_t = tf.concat( (scaling_both, scaling_both), axis=0) optimizer = ScaledLMAOptimizer opt_init_extra_kwargs[ 'diag_mu_scaling_t'] = self._joint_scaling_t if apply_precond: self._joint_precond_t = tf.concat( (scaling_both, scaling_both), axis=0) optimizer = PCGLMAOptimizer opt_init_extra_kwargs[ 'diag_precond_t'] = self._joint_precond_t if apply_diag_mu_scaling: optimizer = ScaledPCGLMAOptimizer # else: optimizer = LMAOptimizer opt_init_args = { "input_var": self.fwd_model.joint_v, "predictions_fn": self._predictions_fn, "loss_fn": self._train_loss_fn, "diag_hessian_fn": self._train_hessian_fn, "name": "opt", "max_cg_iter": max_cg_iter, "min_cg_tol": min_cg_tol, "stochastic_diag_estimator_type": stochastic_diag_estimator_type, "stochastic_diag_estimator_iters": stochastic_diag_estimator_iters, "assert_tolerances": False } opt_init_args.update(opt_init_extra_kwargs) self._attachCustomOptimizerForVariable( optimizer, optimizer_init_args=opt_init_args) self.addCustomMetricToDataLog( title="cg_iters", tensor=self.optimizers[0]._optimizer._total_cg_iterations, log_epoch_frequency=1) self.addCustomMetricToDataLog( title="ls_iters", tensor=self.optimizers[0]._optimizer._total_proj_ls_iterations, log_epoch_frequency=1) self.addCustomMetricToDataLog(title="proj_iters", tensor=self.optimizers[0]._optimizer. _projected_gradient_iterations, log_epoch_frequency=1) self.addCustomMetricToDataLog(title="mu", tensor=self.optimizers[0]._optimizer._mu, log_epoch_frequency=1) self._registration_log_frequency = registration_log_frequency if obj_array_true is not None: self._setObjRegistration( obj_array_true, both_registration_nlse=both_registration_nlse) if probe_wavefront_true is not None: self._setProbeRegistration( probe_wavefront_true, both_registration_nlse=both_registration_nlse) self._addRFactorLog(r_factor_log, registration_log_frequency)
def __init__(self, *args: int, obj_array_true: np.ndarray = None, probe_wavefront_true: np.ndarray = None, obj_abs_proj: bool = True, update_delay_probe: int = 0, update_delay_obj: int = 0, update_delay_aux: int = 0, update_frequency_probe: int = 1, update_frequency_obj: int = 1, update_frequency_aux: int = 1, reconstruct_probe: bool = True, registration_log_frequency: int = 10, r_factor_log: bool = True, loss_type: str = "gaussian", learning_rate_scaling_probe: float = 1.0, learning_rate_scaling_obj: float = 1.0, loss_init_extra_kwargs: dict = None, **kwargs: int): logger.info('initializing...') super().__init__(*args, **kwargs) if self.training_batch_size != self.n_train: e = ValueError("PALM reconstruction does not support minibatch reconstruction.") logger.error(e) raise e if loss_type != "gaussian": e = ValueError("PALM reconstruction does not support other loss functions.") logger.error(e) raise e logger.info('attaching fwd model...') self._attachCustomForwardModel(FarfieldPALMForwardModel, obj_abs_proj=obj_abs_proj) logger.info('creating loss fn...') self._attachCustomLossFunction(AuxLossFunctionT, loss_init_extra_kwargs) logger.info("create learning rates") self._learning_rate_scaling_probe = learning_rate_scaling_probe self._learning_rate_scaling_obj = learning_rate_scaling_obj self._lr_obj = self._getObjLearningRate() * learning_rate_scaling_obj self._lr_probe = self._getProbeLearningRate() * learning_rate_scaling_probe logger.info('creating optimizers...') aux_new_t = self._getAuxUpdate() self._attachCustomOptimizerForVariable(AuxAssignOptimizer, optimizer_init_args={"aux_v": self.fwd_model.aux_v, "aux_new_t": aux_new_t}, initial_update_delay=update_delay_aux, update_frequency=update_frequency_aux) update_delay_obj = update_delay_obj if update_delay_obj is not None else 0 self.attachTensorflowOptimizerForVariable("obj", optimizer_type="gradient", optimizer_init_args={"learning_rate": self._lr_obj}, initial_update_delay=update_delay_obj, update_frequency=update_frequency_obj) update_delay_probe = update_delay_probe if update_delay_probe is not None else 0 if reconstruct_probe: self.attachTensorflowOptimizerForVariable("probe", optimizer_type="gradient", optimizer_init_args={"learning_rate": self._lr_probe}, initial_update_delay=update_delay_probe, update_frequency=update_frequency_probe) if obj_array_true is not None: self.addCustomMetricToDataLog(title="obj_error", tensor=self.fwd_model.obj_cmplx_t, log_epoch_frequency=registration_log_frequency, registration_ground_truth=obj_array_true) if reconstruct_probe and (probe_wavefront_true is not None): self.addCustomMetricToDataLog(title="probe_error", tensor=self.fwd_model.probe_cmplx_t, log_epoch_frequency=registration_log_frequency, registration_ground_truth=probe_wavefront_true) self._addRFactorLog(r_factor_log, 1)
def __init__(self, *args: int, loss_type: str = "least_squared", obj_array_true: np.ndarray = None, probe_wavefront_true: np.ndarray = None, r_factor_log: bool = False, learning_rate_obj: float = 1e-2, update_delay_obj: int = 0, update_delay_probe: int = 0, learning_rate_probe: float = 1e-1, reconstruct_probe: bool = True, registration_log_frequency: int = 10, both_registration_nlse: bool = True, obj_abs_proj: bool = True, loss_init_extra_kwargs: dict = None, **kwargs: int): logger.info('initializing...') super().__init__(*args, **kwargs) logger.info('attaching fwd model...') self.attachForwardModel("farfield", obj_abs_proj=obj_abs_proj) logger.info('creating loss fn...') self.attachLossFunction(loss_type, loss_init_extra_kwargs=loss_init_extra_kwargs) logger.info('creating optimizers...') self.attachTensorflowOptimizerForVariable( "obj", optimizer_type="adam", optimizer_init_args={"learning_rate": learning_rate_obj}, initial_update_delay=update_delay_obj) if reconstruct_probe: self.attachTensorflowOptimizerForVariable( "probe", optimizer_type="adam", optimizer_init_args={"learning_rate": learning_rate_probe}, initial_update_delay=update_delay_probe) if obj_array_true is not None: self.addCustomMetricToDataLog( title="obj_error", tensor=self.fwd_model.obj_cmplx_t, log_epoch_frequency=registration_log_frequency, registration=True, registration_ground_truth=obj_array_true) if both_registration_nlse: self.addCustomMetricToDataLog( title="obj_nlse", tensor=self.fwd_model.obj_cmplx_t, log_epoch_frequency=registration_log_frequency, registration=False, normalized_lse=True, registration_ground_truth=obj_array_true) if reconstruct_probe and (probe_wavefront_true is not None): self.addCustomMetricToDataLog( title="probe_error", tensor=self.fwd_model.probe_cmplx_t, log_epoch_frequency=registration_log_frequency, registration_ground_truth=probe_wavefront_true) if both_registration_nlse: self.addCustomMetricToDataLog( title="probe_nlse", tensor=self.fwd_model.probe_cmplx_t, log_epoch_frequency=registration_log_frequency, registration=False, normalized_lse=True, registration_ground_truth=probe_wavefront_true) self._addRFactorLog(r_factor_log, registration_log_frequency)
def __init__(self, wavelength: float = 1.5e-10, obj: Obj = None, probe_3d: Probe3D = None, scan_grid: BraggPtychoGrid = None, detector: Detector = None, poisson_noise: bool = True, upsampling_factor: int = 1) -> None: self.wavelength = wavelength self.upsampling_factor = upsampling_factor self.poisson_noise = poisson_noise two_theta = scan_grid.two_theta if scan_grid is not None else AngleGridParams( ).two_theta if detector is not None: logger.info("Using supplied detector info.") self.detector = detector else: logger.info("Creating new detector.") self.detector = Detector(**dt.asdict(DetectorParams())) det_3d_shape = (1, *self.detector.shape) obj_xz_nyquist_support = self.wavelength * self.detector.obj_dist / np.array( self.detector.pixel_size) obj_xz_pixel_size = obj_xz_nyquist_support / ( np.array(self.detector.shape) * self.upsampling_factor) # The y pixel size is very ad-hoc obj_y_pixel_size = obj_xz_nyquist_support[1] / self.detector.shape[1] obj_pixel_size = (obj_y_pixel_size, *obj_xz_pixel_size) if obj is not None: logger.info("Using supplied object.") self.obj = obj else: self.obj = self.createObj(dt.asdict(ObjParams()), det_3d_shape, obj_pixel_size, upsampling_factor) probe_xz_shape = np.array(self.detector.shape) * self.upsampling_factor if probe_3d is not None: logger.info("Using supplied 3d probe values.") self.probe_3d = probe_3d # Note that the probe y shape cannot be determined from the other supplied parameters (I think). if (np.any(probe_3d.shape[1:] != probe_xz_shape) or (probe_3d.wavelength != self.wavelength) or np.any(probe_3d.pixel_size != obj_pixel_size)): e = ValueError( f"Mismatch between the supplied probe and the supplied scan parameters." ) logger.error(e) raise e else: self.probe_3d = self.createProbe3D(dt.asdict(Probe2DParams()), wavelength, np.pi / 2 - two_theta, probe_xz_shape, obj_pixel_size) rotate_angle = np.pi / 2 - two_theta probe_y_pix_before_rotation = self.probe_3d.shape[0] * np.cos( rotate_angle) // 1 pady, bordered_obj_ypix = self.calculateObjBorderAfterRotation( rotate_angle, probe_y_pix_before_rotation, self.obj.shape[0], self.obj.shape[1]) if self.obj.bordered_array.shape[0] < bordered_obj_ypix: logger.warning( "Adding zero padding to the object in the y-direction so that the overall object y-width " + "covers the entirety of the feasible probe positions.") self.obj.border_shape = ((pady, pady), self.obj.border_shape[1], self.obj.border_shape[2]) if scan_grid is not None: logger.info("Using supplied scan grid.") self.scan_grid = scan_grid else: elems_yz = lambda tup: np.array(tup)[[0, 2]] scan_grid_params_dict = dt.asdict( Scan2DGridParams(obj_pixel_size=elems_yz(self.obj.pixel_size))) self.scan_grid = self.createScanGrid( scan_grid_params_dict, dt.asdict(AngleGridParams()), elems_yz(self.obj.bordered_array.shape), elems_yz(self.probe_3d.shape)) self._calculateDiffractionPatterns()