""" A simple test of refinement using two crystals. """ from __future__ import annotations from math import pi from cctbx.sgtbx import space_group, space_group_symbols from cctbx.uctbx import unit_cell from dxtbx.model import ScanFactory from dxtbx.model.experiment_list import Experiment, ExperimentList from libtbx.phil import parse from libtbx.test_utils import approx_equal from rstbx.symmetry.constraints.parameter_reduction import symmetrize_reduce_enlarge from scitbx import matrix from scitbx.array_family import flex from dials.algorithms.refinement.parameterisation.beam_parameters import ( BeamParameterisation, ) from dials.algorithms.refinement.parameterisation.crystal_parameters import ( CrystalOrientationParameterisation, CrystalUnitCellParameterisation, ) from dials.algorithms.refinement.parameterisation.detector_parameters import ( DetectorParameterisationSinglePanel, ) from dials.algorithms.refinement.prediction.managed_predictors import ( ScansExperimentsPredictor, ScansRayPredictor, ) from dials.algorithms.refinement.refiner import RefinerFactory, phil_scope from dials.algorithms.spot_prediction import IndexGenerator, ray_intersection from . import geometry_phil, minimiser_phil, setup_geometry def test(args=[]): ############################# # Setup experimental models # ############################# master_phil = parse(f"{geometry_phil}\n{minimiser_phil}") models = setup_geometry.Extract( master_phil, cmdline_args=args, local_overrides="geometry.parameters.random_seed = 1", ) crystal1 = models.crystal models = setup_geometry.Extract( master_phil, cmdline_args=args, local_overrides="geometry.parameters.random_seed = 2", ) mydetector = models.detector mygonio = models.goniometer crystal2 = models.crystal mybeam = models.beam # Build a mock scan for an 18 degree sequence sf = ScanFactory() myscan = sf.make_scan( image_range=(1, 180), exposure_times=0.1, oscillation=(0, 0.1), epochs=list(range(180)), deg=True, ) sequence_range = myscan.get_oscillation_range(deg=False) im_width = myscan.get_oscillation(deg=False)[1] assert sequence_range == (0.0, pi / 10) assert approx_equal(im_width, 0.1 * pi / 180.0) # Build an experiment list experiments = ExperimentList() experiments.append( Experiment( beam=mybeam, detector=mydetector, goniometer=mygonio, scan=myscan, crystal=crystal1, imageset=None, ) ) experiments.append( Experiment( beam=mybeam, detector=mydetector, goniometer=mygonio, scan=myscan, crystal=crystal2, imageset=None, ) ) assert len(experiments.detectors()) == 1 ########################################################## # Parameterise the models (only for perturbing geometry) # ########################################################## det_param = DetectorParameterisationSinglePanel(mydetector) s0_param = BeamParameterisation(mybeam, mygonio) xl1o_param = CrystalOrientationParameterisation(crystal1) xl1uc_param = CrystalUnitCellParameterisation(crystal1) xl2o_param = CrystalOrientationParameterisation(crystal2) xl2uc_param = CrystalUnitCellParameterisation(crystal2) # Fix beam to the X-Z plane (imgCIF geometry), fix wavelength s0_param.set_fixed([True, False, True]) ################################ # Apply known parameter shifts # ################################ # shift detector by 1.0 mm each translation and 2 mrad each rotation det_p_vals = det_param.get_param_vals() p_vals = [a + b for a, b in zip(det_p_vals, [1.0, 1.0, 1.0, 2.0, 2.0, 2.0])] det_param.set_param_vals(p_vals) # shift beam by 2 mrad in free axis s0_p_vals = s0_param.get_param_vals() p_vals = list(s0_p_vals) p_vals[0] += 2.0 s0_param.set_param_vals(p_vals) # rotate crystal a bit (=2 mrad each rotation) xlo_p_vals = [] for xlo in (xl1o_param, xl2o_param): p_vals = xlo.get_param_vals() xlo_p_vals.append(p_vals) new_p_vals = [a + b for a, b in zip(p_vals, [2.0, 2.0, 2.0])] xlo.set_param_vals(new_p_vals) # change unit cell a bit (=0.1 Angstrom length upsets, 0.1 degree of # gamma angle) xluc_p_vals = [] for xluc, xl in ((xl1uc_param, crystal1), ((xl2uc_param, crystal2))): p_vals = xluc.get_param_vals() xluc_p_vals.append(p_vals) cell_params = xl.get_unit_cell().parameters() cell_params = [ a + b for a, b in zip(cell_params, [0.1, 0.1, 0.1, 0.0, 0.0, 0.1]) ] new_uc = unit_cell(cell_params) newB = matrix.sqr(new_uc.fractionalization_matrix()).transpose() S = symmetrize_reduce_enlarge(xl.get_space_group()) S.set_orientation(orientation=newB) X = tuple([e * 1.0e5 for e in S.forward_independent_parameters()]) xluc.set_param_vals(X) ############################# # Generate some reflections # ############################# # All indices in a 2.5 Angstrom sphere for crystal1 resolution = 2.5 index_generator = IndexGenerator( crystal1.get_unit_cell(), space_group(space_group_symbols(1).hall()).type(), resolution, ) indices1 = index_generator.to_array() # All indices in a 2.5 Angstrom sphere for crystal2 resolution = 2.5 index_generator = IndexGenerator( crystal2.get_unit_cell(), space_group(space_group_symbols(1).hall()).type(), resolution, ) indices2 = index_generator.to_array() # Predict rays within the sequence range. Set experiment IDs ray_predictor = ScansRayPredictor(experiments, sequence_range) obs_refs1 = ray_predictor(indices1, experiment_id=0) obs_refs1["id"] = flex.int(len(obs_refs1), 0) obs_refs2 = ray_predictor(indices2, experiment_id=1) obs_refs2["id"] = flex.int(len(obs_refs2), 1) # Take only those rays that intersect the detector intersects = ray_intersection(mydetector, obs_refs1) obs_refs1 = obs_refs1.select(intersects) intersects = ray_intersection(mydetector, obs_refs2) obs_refs2 = obs_refs2.select(intersects) # Make a reflection predictor and re-predict for all these reflections. The # result is the same, but we gain also the flags and xyzcal.px columns ref_predictor = ScansExperimentsPredictor(experiments) obs_refs1 = ref_predictor(obs_refs1) obs_refs2 = ref_predictor(obs_refs2) # Set 'observed' centroids from the predicted ones obs_refs1["xyzobs.mm.value"] = obs_refs1["xyzcal.mm"] obs_refs2["xyzobs.mm.value"] = obs_refs2["xyzcal.mm"] # Invent some variances for the centroid positions of the simulated data im_width = 0.1 * pi / 18.0 px_size = mydetector[0].get_pixel_size() var_x = flex.double(len(obs_refs1), (px_size[0] / 2.0) ** 2) var_y = flex.double(len(obs_refs1), (px_size[1] / 2.0) ** 2) var_phi = flex.double(len(obs_refs1), (im_width / 2.0) ** 2) obs_refs1["xyzobs.mm.variance"] = flex.vec3_double(var_x, var_y, var_phi) var_x = flex.double(len(obs_refs2), (px_size[0] / 2.0) ** 2) var_y = flex.double(len(obs_refs2), (px_size[1] / 2.0) ** 2) var_phi = flex.double(len(obs_refs2), (im_width / 2.0) ** 2) obs_refs2["xyzobs.mm.variance"] = flex.vec3_double(var_x, var_y, var_phi) # concatenate reflection lists obs_refs1.extend(obs_refs2) obs_refs = obs_refs1 ############################### # Undo known parameter shifts # ############################### s0_param.set_param_vals(s0_p_vals) det_param.set_param_vals(det_p_vals) xl1o_param.set_param_vals(xlo_p_vals[0]) xl2o_param.set_param_vals(xlo_p_vals[1]) xl1uc_param.set_param_vals(xluc_p_vals[0]) xl2uc_param.set_param_vals(xluc_p_vals[1]) # scan static first params = phil_scope.fetch(source=parse("")).extract() refiner = RefinerFactory.from_parameters_data_experiments( params, obs_refs, experiments ) refiner.run() # scan varying params.refinement.parameterisation.scan_varying = True refiner = RefinerFactory.from_parameters_data_experiments( params, obs_refs, experiments ) refiner.run() # Ensure all models have scan-varying state set # (https://github.com/dials/dials/issues/798) refined_experiments = refiner.get_experiments() sp = [xl.get_num_scan_points() for xl in refined_experiments.crystals()] assert sp.count(181) == 2