def index(self): experiments = ExperimentList() had_refinement_error = False have_similar_crystal_models = False while True: if had_refinement_error or have_similar_crystal_models: break max_lattices = self.params.multiple_lattice_search.max_lattices if max_lattices is not None and len(experiments) >= max_lattices: break if len(experiments) > 0: cutoff_fraction = (self.params.multiple_lattice_search. recycle_unindexed_reflections_cutoff) d_spacings = 1 / self.reflections["rlp"].norms() d_min_indexed = flex.min( d_spacings.select(self.indexed_reflections)) min_reflections_for_indexing = cutoff_fraction * len( self.reflections.select(d_spacings > d_min_indexed)) crystal_ids = self.reflections.select( d_spacings > d_min_indexed)["id"] if (crystal_ids == -1).count(True) < min_reflections_for_indexing: logger.info( "Finish searching for more lattices: %i unindexed reflections remaining." % ((crystal_ids == -1).count(True))) break n_lattices_previous_cycle = len(experiments) if self.d_min is None: self.d_min = self.params.refinement_protocol.d_min_start if len(experiments) == 0: new_expts = self.find_lattices() generate_experiment_identifiers(new_expts) experiments.extend(new_expts) else: try: new = self.find_lattices() generate_experiment_identifiers(new) experiments.extend(new) except DialsIndexError: logger.info("Indexing remaining reflections failed") if self.params.refinement_protocol.d_min_step is libtbx.Auto: n_cycles = self.params.refinement_protocol.n_macro_cycles if self.d_min is None or n_cycles == 1: self.params.refinement_protocol.d_min_step = 0 else: d_spacings = 1 / self.reflections["rlp"].norms() d_min_all = flex.min(d_spacings) self.params.refinement_protocol.d_min_step = ( self.d_min - d_min_all) / (n_cycles - 1) logger.info("Using d_min_step %.1f" % self.params.refinement_protocol.d_min_step) if len(experiments) == 0: raise DialsIndexError("No suitable lattice could be found.") elif len(experiments) == n_lattices_previous_cycle: # no more lattices found break for i_cycle in range( self.params.refinement_protocol.n_macro_cycles): if (i_cycle > 0 and self.d_min is not None and self.params.refinement_protocol.d_min_step > 0): d_min = self.d_min - self.params.refinement_protocol.d_min_step d_min = max(d_min, 0) if self.params.refinement_protocol.d_min_final is not None: d_min = max( d_min, self.params.refinement_protocol.d_min_final) if d_min >= 0: self.d_min = d_min logger.info("Increasing resolution to %.2f Angstrom" % d_min) # reset reflection lattice flags # the lattice a given reflection belongs to: a value of -1 indicates # that a reflection doesn't belong to any lattice so far self.reflections["id"] = flex.int(len(self.reflections), -1) self.index_reflections(experiments, self.reflections) if i_cycle == 0 and self.params.known_symmetry.space_group is not None: self._apply_symmetry_post_indexing( experiments, self.reflections, n_lattices_previous_cycle) logger.info("\nIndexed crystal models:") self.show_experiments(experiments, self.reflections, d_min=self.d_min) if self._check_have_similar_crystal_models(experiments): have_similar_crystal_models = True break logger.info("") logger.info("#" * 80) logger.info("Starting refinement (macro-cycle %i)" % (i_cycle + 1)) logger.info("#" * 80) logger.info("") self.indexed_reflections = self.reflections["id"] > -1 sel = flex.bool(len(self.reflections), False) lengths = 1 / self.reflections["rlp"].norms() if self.d_min is not None: isel = (lengths <= self.d_min).iselection() sel.set_selected(isel, True) sel.set_selected(self.reflections["id"] == -1, True) self.reflections.unset_flags(sel, self.reflections.flags.indexed) self.unindexed_reflections = self.reflections.select(sel) reflections_for_refinement = self.reflections.select( self.indexed_reflections) if self.params.refinement_protocol.mode == "repredict_only": refined_experiments, refined_reflections = ( experiments, reflections_for_refinement, ) from dials.algorithms.refinement.prediction.managed_predictors import ( ExperimentsPredictorFactory, ) ref_predictor = ExperimentsPredictorFactory.from_experiments( experiments, spherical_relp=self.all_params.refinement. parameterisation.spherical_relp_model, ) ref_predictor(refined_reflections) else: try: refined_experiments, refined_reflections = self.refine( experiments, reflections_for_refinement) except (DialsRefineConfigError, DialsRefineRuntimeError) as e: if len(experiments) == 1: raise DialsIndexRefineError(str(e)) had_refinement_error = True logger.info("Refinement failed:") logger.info(e) del experiments[-1] # remove experiment id from the reflections associated # with this deleted experiment - indexed flag removed # below last = len(experiments) sel = refined_reflections["id"] == last logger.info("Removing %d reflections with id %d" % (sel.count(True), last)) refined_reflections["id"].set_selected(sel, -1) break self._unit_cell_volume_sanity_check(experiments, refined_experiments) self.refined_reflections = refined_reflections self.refined_reflections.unset_flags( self.refined_reflections["id"] < 0, self.refined_reflections.flags.indexed, ) for i, expt in enumerate(self.experiments): ref_sel = self.refined_reflections.select( self.refined_reflections["imageset_id"] == i) ref_sel = ref_sel.select(ref_sel["id"] >= 0) for i_expt in set(ref_sel["id"]): refined_expt = refined_experiments[i_expt] expt.detector = refined_expt.detector expt.beam = refined_expt.beam expt.goniometer = refined_expt.goniometer expt.scan = refined_expt.scan refined_expt.imageset = expt.imageset if not (self.all_params.refinement.parameterisation.beam.fix == "all" and self.all_params.refinement. parameterisation.detector.fix == "all"): # Experimental geometry may have changed - re-map centroids to # reciprocal space self.reflections.map_centroids_to_reciprocal_space( self.experiments) # update for next cycle experiments = refined_experiments self.refined_experiments = refined_experiments logger.info("\nRefined crystal models:") self.show_experiments(self.refined_experiments, self.reflections, d_min=self.d_min) if (i_cycle >= 2 and self.d_min == self.params.refinement_protocol.d_min_final): logger.info( "Target d_min_final reached: finished with refinement") break if self.refined_experiments is None: raise DialsIndexRefineError( "None of the experiments could refine.") if len(self.refined_experiments) > 1: from dials.algorithms.indexing.compare_orientation_matrices import ( rotation_matrix_differences, ) logger.info( rotation_matrix_differences( self.refined_experiments.crystals())) self._xyzcal_mm_to_px(self.refined_experiments, self.refined_reflections)
def index(self): # most of this is the same as dials.algorithms.indexing.indexer.indexer_base.index(), with some stills # specific modifications (don't re-index after choose best orientation matrix, but use the indexing from # choose best orientation matrix, also don't use macrocycles) of refinement after indexing. # 2017 update: do accept multiple lattices per shot experiments = ExperimentList() while True: self.d_min = self.params.refinement_protocol.d_min_start max_lattices = self.params.multiple_lattice_search.max_lattices if max_lattices is not None and len(experiments) >= max_lattices: break if len(experiments) > 0: cutoff_fraction = ( self.params.multiple_lattice_search.recycle_unindexed_reflections_cutoff ) d_spacings = 1 / self.reflections["rlp"].norms() d_min_indexed = flex.min(d_spacings.select(self.indexed_reflections)) min_reflections_for_indexing = cutoff_fraction * len( self.reflections.select(d_spacings > d_min_indexed) ) crystal_ids = self.reflections.select(d_spacings > d_min_indexed)["id"] if (crystal_ids == -1).count(True) < min_reflections_for_indexing: logger.info( "Finish searching for more lattices: %i unindexed reflections remaining." % (min_reflections_for_indexing) ) break n_lattices_previous_cycle = len(experiments) # index multiple lattices per shot if len(experiments) == 0: new = self.find_lattices() generate_experiment_identifiers(new) experiments.extend(new) if len(experiments) == 0: raise DialsIndexError("No suitable lattice could be found.") else: try: new = self.find_lattices() generate_experiment_identifiers(new) experiments.extend(new) except Exception as e: logger.info("Indexing remaining reflections failed") logger.debug( "Indexing remaining reflections failed, exception:\n" + str(e) ) # reset reflection lattice flags # the lattice a given reflection belongs to: a value of -1 indicates # that a reflection doesn't belong to any lattice so far self.reflections["id"] = flex.int(len(self.reflections), -1) self.index_reflections(experiments, self.reflections) if len(experiments) == n_lattices_previous_cycle: # no more lattices found break if ( not self.params.stills.refine_candidates_with_known_symmetry and self.params.known_symmetry.space_group is not None ): self._apply_symmetry_post_indexing( experiments, self.reflections, n_lattices_previous_cycle ) # discard nearly overlapping lattices on the same shot if self._check_have_similar_crystal_models(experiments): break self.indexed_reflections = self.reflections["id"] > -1 if self.d_min is None: sel = self.reflections["id"] <= -1 else: sel = flex.bool(len(self.reflections), False) lengths = 1 / self.reflections["rlp"].norms() isel = (lengths >= self.d_min).iselection() sel.set_selected(isel, True) sel.set_selected(self.reflections["id"] > -1, False) self.unindexed_reflections = self.reflections.select(sel) reflections_for_refinement = self.reflections.select( self.indexed_reflections ) if len(self.params.stills.isoforms) > 0: logger.info("") logger.info("#" * 80) logger.info("Starting refinement") logger.info("#" * 80) logger.info("") isoform_experiments = ExperimentList() isoform_reflections = flex.reflection_table() # Note, changes to params after initial indexing. Cannot use tie to target when fixing the unit cell. self.all_params.refinement.reflections.outlier.algorithm = "null" self.all_params.refinement.parameterisation.crystal.fix = "cell" self.all_params.refinement.parameterisation.crystal.unit_cell.restraints.tie_to_target = ( [] ) for expt_id, experiment in enumerate(experiments): reflections = reflections_for_refinement.select( reflections_for_refinement["id"] == expt_id ) reflections["id"] = flex.int(len(reflections), 0) refiners = [] for isoform in self.params.stills.isoforms: iso_experiment = copy.deepcopy(experiment) crystal = iso_experiment.crystal if ( isoform.lookup_symbol != crystal.get_space_group().type().lookup_symbol() ): logger.info( "Crystal isoform lookup_symbol %s does not match isoform %s lookup_symbol %s" % ( crystal.get_space_group().type().lookup_symbol(), isoform.name, isoform.lookup_symbol, ) ) continue crystal.set_B(isoform.cell.fractionalization_matrix()) logger.info("Refining isoform %s" % isoform.name) refiners.append( e_refine( params=self.all_params, experiments=ExperimentList([iso_experiment]), reflections=reflections, graph_verbose=False, ) ) if len(refiners) == 0: raise DialsIndexError( "No isoforms had a lookup symbol that matched" ) positional_rmsds = [ math.sqrt(P.rmsds()[0] ** 2 + P.rmsds()[1] ** 2) for P in refiners ] logger.info( "Positional rmsds for all isoforms:" + str(positional_rmsds) ) minrmsd_mm = min(positional_rmsds) minindex = positional_rmsds.index(minrmsd_mm) logger.info( "The smallest rmsd is %5.1f um from isoform %s" % ( 1000.0 * minrmsd_mm, self.params.stills.isoforms[minindex].name, ) ) if self.params.stills.isoforms[minindex].rmsd_target_mm is not None: logger.info( "Asserting %f < %f" % ( minrmsd_mm, self.params.stills.isoforms[minindex].rmsd_target_mm, ) ) assert ( minrmsd_mm < self.params.stills.isoforms[minindex].rmsd_target_mm ) logger.info( "Acceptable rmsd for isoform %s." % (self.params.stills.isoforms[minindex].name) ) if len(self.params.stills.isoforms) == 2: logger.info( "Rmsd gain over the other isoform %5.1f um." % (1000.0 * abs(positional_rmsds[0] - positional_rmsds[1])) ) R = refiners[minindex] # Now one last check to see if direct beam is out of bounds if self.params.stills.isoforms[minindex].beam_restraint is not None: from scitbx import matrix refined_beam = matrix.col( R.get_experiments()[0] .detector[0] .get_beam_centre_lab(experiments[0].beam.get_s0())[0:2] ) known_beam = matrix.col( self.params.stills.isoforms[minindex].beam_restraint ) logger.info( "Asserting difference in refined beam center and expected beam center %f < %f" % ( (refined_beam - known_beam).length(), self.params.stills.isoforms[minindex].rmsd_target_mm, ) ) assert ( (refined_beam - known_beam).length() < self.params.stills.isoforms[minindex].rmsd_target_mm ) # future--circle of confusion could be given as a separate length in mm instead of reusing rmsd_target experiment = R.get_experiments()[0] experiment.crystal.identified_isoform = self.params.stills.isoforms[ minindex ].name isoform_experiments.append(experiment) reflections["id"] = flex.int(len(reflections), expt_id) isoform_reflections.extend(reflections) experiments = isoform_experiments reflections_for_refinement = isoform_reflections if self.params.refinement_protocol.mode == "repredict_only": from dials.algorithms.indexing.nave_parameters import NaveParameters from dials.algorithms.refinement.prediction.managed_predictors import ( ExperimentsPredictorFactory, ) refined_experiments, refined_reflections = ( experiments, reflections_for_refinement, ) ref_predictor = ExperimentsPredictorFactory.from_experiments( experiments, force_stills=True, spherical_relp=self.all_params.refinement.parameterisation.spherical_relp_model, ) ref_predictor(refined_reflections) refined_reflections["delpsical2"] = ( refined_reflections["delpsical.rad"] ** 2 ) for expt_id in range(len(refined_experiments)): refls = refined_reflections.select( refined_reflections["id"] == expt_id ) nv = NaveParameters( params=self.all_params, experiments=refined_experiments[expt_id : expt_id + 1], reflections=refls, refinery=None, graph_verbose=False, ) experiments[expt_id].crystal = nv() ref_predictor = ExperimentsPredictorFactory.from_experiments( experiments, force_stills=True, spherical_relp=self.all_params.refinement.parameterisation.spherical_relp_model, ) ref_predictor(refined_reflections) elif self.params.refinement_protocol.mode is None: refined_experiments, refined_reflections = ( experiments, reflections_for_refinement, ) else: try: refined_experiments, refined_reflections = self.refine( experiments, reflections_for_refinement ) except Exception as e: s = str(e) if len(experiments) == 1: raise DialsIndexRefineError(e.message) logger.info("Refinement failed:") logger.info(s) del experiments[-1] break self._unit_cell_volume_sanity_check(experiments, refined_experiments) self.refined_reflections = refined_reflections.select( refined_reflections["id"] > -1 ) for i, expt in enumerate(self.experiments): ref_sel = self.refined_reflections.select( self.refined_reflections["imageset_id"] == i ) ref_sel = ref_sel.select(ref_sel["id"] >= 0) for i_expt in set(ref_sel["id"]): refined_expt = refined_experiments[i_expt] expt.detector = refined_expt.detector expt.beam = refined_expt.beam expt.goniometer = refined_expt.goniometer expt.scan = refined_expt.scan refined_expt.imageset = expt.imageset if not ( self.all_params.refinement.parameterisation.beam.fix == "all" and self.all_params.refinement.parameterisation.detector.fix == "all" ): # Experimental geometry may have changed - re-map centroids to # reciprocal space self.reflections.map_centroids_to_reciprocal_space(self.experiments) # update for next cycle experiments = refined_experiments self.refined_experiments = refined_experiments if self.refined_experiments is None: raise DialsIndexRefineError("None of the experiments could refine.") # discard experiments with zero reflections after refinement id_set = set(self.refined_reflections["id"]) if len(id_set) < len(self.refined_experiments): filtered_refined_reflections = flex.reflection_table() for i in range(len(self.refined_experiments)): if i not in id_set: del self.refined_experiments[i] for old, new in zip(sorted(id_set), range(len(id_set))): subset = self.refined_reflections.select( self.refined_reflections["id"] == old ) subset["id"] = flex.int(len(subset), new) filtered_refined_reflections.extend(subset) self.refined_reflections = filtered_refined_reflections if len(self.refined_experiments) > 1: from dials.algorithms.indexing.compare_orientation_matrices import ( rotation_matrix_differences, ) logger.info( rotation_matrix_differences(self.refined_experiments.crystals()) ) logger.info("Final refined crystal models:") for i, crystal_model in enumerate(self.refined_experiments.crystals()): n_indexed = 0 for _ in experiments.where(crystal=crystal_model): n_indexed += (self.reflections["id"] == i).count(True) logger.info("model %i (%i reflections):" % (i + 1, n_indexed)) logger.info(crystal_model) if ( "xyzcal.mm" in self.refined_reflections ): # won't be there if refine_all_candidates = False and no isoforms self._xyzcal_mm_to_px(self.experiments, self.refined_reflections)
def test_compare_orientation_matrices(): # try and see if we can get back the original rotation matrix and euler angles real_space_a = matrix.col((10, 0, 0)) real_space_b = matrix.col((0, 10, 10)) real_space_c = matrix.col((0, 0, 10)) euler_angles = (1.3, 5.6, 7.8) R = matrix.sqr( euler.xyz_matrix(euler_angles[0], euler_angles[1], euler_angles[2])) crystal_a = Crystal(real_space_a, real_space_b, real_space_c, space_group=sgtbx.space_group("P 1")) crystal_b = Crystal( R * real_space_a, R * real_space_b, R * real_space_c, space_group=sgtbx.space_group("P 1"), ) assert (matrix.sqr(crystal_b.get_U()) * matrix.sqr(crystal_a.get_U()).transpose()).elems == pytest.approx( R.elems) ( best_R_ab, best_axis, best_angle, best_cb_op, ) = compare_orientation_matrices.difference_rotation_matrix_axis_angle( crystal_a, crystal_b) best_euler_angles = euler.xyz_angles(best_R_ab) assert best_euler_angles == pytest.approx(euler_angles) assert best_cb_op.is_identity_op() assert best_R_ab.elems == pytest.approx(R.elems) # now see if we can deconvolute the original euler angles after applying # a change of basis to one of the crystals crystal_a = Crystal(real_space_a, real_space_b, real_space_c, space_group=sgtbx.space_group("I 2 3")) cb_op = sgtbx.change_of_basis_op("z,x,y") crystal_b = Crystal( R * real_space_a, R * real_space_b, R * real_space_c, space_group=sgtbx.space_group("I 2 3"), ).change_basis(cb_op) ( best_R_ab, best_axis, best_angle, best_cb_op, ) = compare_orientation_matrices.difference_rotation_matrix_axis_angle( crystal_a, crystal_b) best_euler_angles = euler.xyz_angles(best_R_ab) assert best_euler_angles == pytest.approx(euler_angles) assert best_cb_op.c() == cb_op.inverse().c() assert best_R_ab.elems == pytest.approx(R.elems) crystal_c = crystal_b.change_basis(sgtbx.change_of_basis_op("-y,-z,x")) assert crystal_c != crystal_b s = compare_orientation_matrices.rotation_matrix_differences( [crystal_a, crystal_b, crystal_c], comparison="pairwise") s = "\n".join(s.splitlines()[:-1]).replace("-0.000", "0.000") print(s) assert (s == """\ Change of basis op: b,c,a Rotation matrix to transform crystal 1 to crystal 2: {{0.986, -0.135, 0.098}, {0.138, 0.990, -0.023}, {-0.094, 0.036, 0.995}} Rotation of 9.738 degrees about axis (0.172, 0.565, 0.807) Change of basis op: -a,-b,c Rotation matrix to transform crystal 1 to crystal 3: {{0.986, -0.135, 0.098}, {0.138, 0.990, -0.023}, {-0.094, 0.036, 0.995}} Rotation of 9.738 degrees about axis (0.172, 0.565, 0.807) Change of basis op: c,-a,-b Rotation matrix to transform crystal 2 to crystal 3: {{1.000, 0.000, 0.000}, {0.000, 1.000, 0.000}, {0.000, 0.000, 1.000}}""") s = compare_orientation_matrices.rotation_matrix_differences( [crystal_a, crystal_b, crystal_c], comparison="sequential") s = "\n".join(s.splitlines()[:-1]).replace("-0.000", "0.000") print(s) assert (s == """\ Change of basis op: b,c,a Rotation matrix to transform crystal 1 to crystal 2: {{0.986, -0.135, 0.098}, {0.138, 0.990, -0.023}, {-0.094, 0.036, 0.995}} Rotation of 9.738 degrees about axis (0.172, 0.565, 0.807) Change of basis op: c,-a,-b Rotation matrix to transform crystal 2 to crystal 3: {{1.000, 0.000, 0.000}, {0.000, 1.000, 0.000}, {0.000, 0.000, 1.000}}""") s = compare_orientation_matrices.rotation_matrix_differences( (crystal_a, crystal_b), miller_indices=((1, 0, 0), (1, 1, 0))) assert (s == """\ Change of basis op: b,c,a Rotation matrix to transform crystal 1 to crystal 2: {{0.986, -0.135, 0.098}, {0.138, 0.990, -0.023}, {-0.094, 0.036, 0.995}} Rotation of 9.738 degrees about axis (0.172, 0.565, 0.807) (1,0,0): 15.26 deg (1,1,0): 9.12 deg """)
def main(rank): device_Id = rank % ngpu worker_Id = node_id * ngpu + rank import os import sys from copy import deepcopy import glob from itertools import izip from scipy.spatial import distance import h5py import scipy.ndimage from IPython import embed import numpy as np import pandas from scipy.spatial import cKDTree from simtbx.nanoBragg import shapetype, nanoBragg from libtbx.phil import parse from scitbx.matrix import sqr import dxtbx from dxtbx.model.experiment_list import ExperimentListFactory from dxtbx.model.crystal import CrystalFactory from dials.algorithms.indexing.compare_orientation_matrices \ import rotation_matrix_differences from dials.array_family import flex from dials.command_line.find_spots import phil_scope as find_spots_phil_scope from cxid9114.refine import metrics from cxid9114 import utils from cxid9114.geom import geom_utils from cxid9114.spots import integrate, spot_utils from cxid9114 import parameters from cxid9114.sim import sim_utils from cctbx import miller, sgtbx from cxid9114 import utils from cxid9114.bigsim import sim_spectra from cxid9114.refine.jitter_refine import make_param_list spot_par = find_spots_phil_scope.fetch(source=parse("")).extract() spot_par.spotfinder.threshold.dispersion.global_threshold = 40 spot_par.spotfinder.threshold.dispersion.gain = 28 spot_par.spotfinder.threshold.dispersion.kernel_size = [2, 2] spot_par.spotfinder.threshold.dispersion.sigma_strong = 1 spot_par.spotfinder.threshold.dispersion.sigma_background = 6 spot_par.spotfinder.filter.min_spot_size = 3 spot_par.spotfinder.force_2d = True odir = args.odir odirj = os.path.join(odir, "job%d" % worker_Id) #all_pkl_files = [s for sl in \ # [ files for _,_, files in os.walk(odir)]\ # for s in sl if s.endswith("pkl")] #print "Found %d pkl files already in %s!" \ # % (len(all_pkl_files), odir) if not os.path.exists(odirj): os.makedirs(odirj) hkl_tol = .15 run = 61 shot_idx = 0 ENERGIES = [parameters.ENERGY_LOW, parameters.ENERGY_HIGH] # colors of the beams FF = [10000, None] cryst_descr = { '__id__': 'crystal', 'real_space_a': (79, 0, 0), 'real_space_b': (0, 79, 0), 'real_space_c': (0, 0, 38), 'space_group_hall_symbol': '-P 4 2' } crystalAB = CrystalFactory.from_dict(cryst_descr) sfall_main = sim_spectra.load_spectra("../bigsim/test_sfall.h5") FFdat = [sfall_main[19], sfall_main[110]] FLUX = [1e11, 1e11] # fluxes of the beams chanA_flux = 1e11 chanB_flux = 1e11 FLUXdat = [chanA_flux, chanB_flux] GAIN = 1 waveA = parameters.ENERGY_CONV / ENERGIES[0] waveB = parameters.ENERGY_CONV / ENERGIES[1] from cxid9114.bigsim.bigsim_geom import DET, BEAM detector = DET print("Rank %d Begin" % worker_Id) for i_data in range(args.num_trials): pklname = "%s_rank%d_data%d.pkl" % (ofile, worker_Id, i_data) pklname = os.path.join(odirj, pklname) print("<><><><><><><") print("Job %d: trial %d / %d" % (worker_Id, i_data + 1, args.num_trials)) print("<><><><><><><") if (worker_Id == 0 and i_data % smi_stride == 0 and cuda): print("GPU status") os.system("nvidia-smi") print("\n\n") print("CPU memory usage") mem_usg = """ps -U dermen --no-headers -o rss | awk '{ sum+=$1} END {print int(sum/1024) "MB consumed by CPU user"}'""" os.system(mem_usg) beamA = deepcopy(BEAM) beamB = deepcopy(BEAM) beamA.set_wavelength(waveA) beamB.set_wavelength(waveB) np.random.seed(args.seed) crystalAB = CrystalFactory.from_dict(cryst_descr) randnums = np.random.random(3) Rrand = random_rotation(1, randnums) crystalAB.set_U(Rrand.ravel()) #pert = np.random.uniform(0.0001/2/np.pi, 0.0003 / 2. /np.pi) #print("PERT %f" % pert) #Rsmall = random_rotation(0.00001, randnums ) #pert) params_lst = make_param_list(crystalAB, DET, BEAM, 1, rot=0.08, cell=.0000001, eq=(1, 1, 0), min_Ncell=23, max_Ncell=24, min_mos_spread=0.02, max_mos_spread=0.08) Ctruth = params_lst[0]['crystal'] print Ctruth.get_unit_cell().parameters() print crystalAB.get_unit_cell().parameters() init_comp = rotation_matrix_differences((Ctruth, crystalAB)) init_rot = float(init_comp.split("\n")[-2].split()[2]) if use_data_spec: print "NOT IMPLEMENTED, Using a phony 2col spectrum to simulate the data" data_fluxes = FLUXdat data_energies = [parameters.ENERGY_LOW, parameters.ENERGY_HIGH] data_ff = FFdat else: print "Using a phony two color spectrum to simulate the data" data_fluxes = FLUXdat data_energies = [parameters.ENERGY_LOW, parameters.ENERGY_HIGH] data_ff = FFdat print("Truth crystal Misorientation deviation: %f deg" % init_rot) if args.truth_cryst: print "Using truth crystal" dataCryst = Ctruth else: print "Not using truth crystal" dataCryst = crystalAB if not make_background: print "SIMULATING Flat-Fhkl IMAGES" simsAB = sim_utils.sim_twocolors2( crystalAB, detector, BEAM, FF, [parameters.ENERGY_LOW, parameters.ENERGY_HIGH], FLUX, pids=None, Gauss=Gauss, cuda=cuda, oversample=oversample, Ncells_abc=Ncells_abc, mos_dom=mos_doms, mos_spread=mos_spread, exposure_s=exposure_s, beamsize_mm=beamsize_mm, device_Id=device_Id, boost=boost) if make_background: print("MAKING BACKGROUND") spec_file = h5py.File("../bigsim/simMe_data_run62.h5", "r") ave_spec = np.mean(spec_file["hist_spec"][()], axis=0) data_fluxes = [ave_spec[19], ave_spec[110]] data_energies = spec_file["energy_bins"][()][[19, 110]] data_ff = [1, 1] #*len(data_energies) only_water = True else: only_water = False print "SIULATING DATA IMAGE" print data_fluxes simsDataSum = sim_utils.sim_twocolors2(dataCryst, detector, BEAM, data_ff, data_energies, data_fluxes, pids=None, Gauss=Gauss, cuda=cuda, oversample=oversample, Ncells_abc=Ncells_abc, accumulate=True, mos_dom=mos_doms, mos_spread=mos_spread, boost=boost, exposure_s=exposure_s, beamsize_mm=beamsize_mm, only_water=only_water, device_Id=device_Id) simsDataSum = np.array(simsDataSum) if make_background: bg_out = h5py.File(bg_name, "w") bg_out.create_dataset("bigsim_d9114", data=simsDataSum[0]) print "Background made! Saved to file %s" % bg_name sys.exit() if add_background: print("ADDING BG") background = h5py.File(bg_name, "r")['bigsim_d9114'][()] bg_scale = np.sum([39152412349.12075, 32315440627.406036]) bg_scale = np.sum(data_fluxes) / bg_scale print "%.3e backgorund scale" % bg_scale print "BG shape", background.shape simsDataSum[0] += background * bg_scale if add_noise: print("ADDING NOISE") for pidx in range(1): SIM = nanoBragg(detector=DET, beam=BEAM, panel_id=pidx) SIM.exposure_s = exposure_s SIM.beamsize_mm = beamsize_mm SIM.flux = np.sum(data_fluxes) SIM.detector_psf_kernel_radius_pixels = 5 SIM.detector_psf_type = shapetype.Unknown # for CSPAD SIM.detector_psf_fwhm_mm = 0 SIM.quantum_gain = 28 SIM.raw_pixels = flex.double(simsDataSum[pidx].ravel()) SIM.add_noise() simsDataSum[pidx] = SIM.raw_pixels.as_numpy_array()\ .reshape(simsDataSum[0].shape) SIM.free_all() del SIM if args.write_img: print "SAVING DATAFILE" h5name = "%s_rank%d_data%d.h5" % (ofile, worker_Id, i_data) h5name = os.path.join(odirj, h5name) fout = h5py.File(h5name, "w") fout.create_dataset("bigsim_d9114", data=simsDataSum[0]) fout.create_dataset("crystalAB", data=crystalAB.get_A()) fout.create_dataset("dataCryst", data=dataCryst.get_A()) fout.close() if args.write_sim_img: print "SAVING DATAFILE" for i_sim in simsAB: sim_h5name = "%s_rank%d_sim%d_%d.h5" % (ofile, worker_Id, i_data, i_sim) sim_h5name = os.path.join(odirj, sim_h5name) from IPython import embed embed() fout = h5py.File(sim_h5name, "w") fout.create_dataset("bigsim_d9114", data=simsAB[i_sim][0]) fout.create_dataset("crystalAB", data=crystalAB.get_A()) fout.create_dataset("dataCryst", data=dataCryst.get_A()) fout.close() print "RELFS FROM SIMS" refl_simA = spot_utils.refls_from_sims(simsAB[0], detector, beamA, thresh=thresh) refl_simB = spot_utils.refls_from_sims(simsAB[1], detector, beamB, thresh=thresh) if use_dials_spotter: print("DIALS SPOTTING") El = utils.explist_from_numpyarrays(simsDataSum, DET, beamA) refl_data = flex.reflection_table.from_observations(El, spot_par) print("Found %d refls using DIALS spot finder" % len(refl_data)) else: refl_data = spot_utils.refls_from_sims(simsDataSum, detector, beamA,\ thresh=thresh) print("Found %d refls using threshold" % len(refl_data)) if len(refl_data) == 0: print "Rank %d: No reflections found! " % (worker_Id) continue residA = metrics.check_indexable2(refl_data, refl_simA, detector, beamA, crystalAB, hkl_tol) residB = metrics.check_indexable2(refl_data, refl_simB, detector, beamB, crystalAB, hkl_tol) sg96 = sgtbx.space_group(" P 4nw 2abw") FA = sfall_main[19] # utils.open_flex('SA.pkl') # ground truth values FB = sfall_main[110] #utils.open_flex('SB.pkl') # ground truth values HA = tuple([hkl for hkl in FA.indices()]) HB = tuple([hkl for hkl in FB.indices()]) HA_val_map = {h: data for h, data in izip(FA.indices(), FA.data())} HB_val_map = {h: data for h, data in izip(FB.indices(), FB.data())} def get_val_at_hkl(hkl, val_map): poss_equivs = [ i.h() for i in miller.sym_equiv_indices(sg96, hkl).indices() ] in_map = False for hkl2 in poss_equivs: if hkl2 in val_map: # fast lookup in_map = True break if in_map: return hkl2, val_map[hkl2] else: return (None, None, None), -1 filt = 1 #True #`False #True if filt: _, all_HiA, _ = spot_utils.refls_to_hkl(refl_simA, detector, beamA, crystal=crystalAB, returnQ=True) all_treeA = cKDTree(all_HiA) nnA = all_treeA.query_ball_point(all_HiA, r=1e-7) _, all_HiB, _ = spot_utils.refls_to_hkl(refl_simB, detector, beamB, crystal=crystalAB, returnQ=True) all_treeB = cKDTree(all_HiB) nnB = all_treeB.query_ball_point(all_HiB, r=1e-7) NreflA = len(refl_simA) NreflB = len(refl_simB) drop_meA = [] for i, vals in enumerate(nnA): if i in drop_meA: continue if len(vals) > 1: pids = [refl_simA[v]['panel'] for v in vals] if len(set(pids)) == 1: refl_vals = refl_simA.select( flex.bool( [i_v in vals for i_v in np.arange(NreflA)])) x, y, z = spot_utils.xyz_from_refl(refl_vals) allI = [r['intensity.sum.value'] for r in refl_vals] allI = sum(allI) xm = np.mean(x) ym = np.mean(y) zm = np.mean(z) drop_meA.extend(vals[1:]) x1b, x2b, y1b, y2b, z1b, z2b = zip( *[r['bbox'] for r in refl_vals]) keep_me = vals[0] # indexing order is important to modify as reference refl_simA['intensity.sum.value'][keep_me] = allI refl_simA['xyzobs.px.value'][keep_me] = (xm, ym, zm) refl_simA['bbox'][keep_me] = (min(x1b), max(x2b),\ min(y1b), max(y2b), min(z1b), max(z2b)) else: drop_meA.append(vals) print vals if drop_meA: keep_meA = np.array([i not in drop_meA for i in range(NreflA)]) refl_simA = refl_simA.select(flex.bool(keep_meA)) NreflA = len(refl_simA) drop_meB = [] for i, vals in enumerate(nnB): if i in drop_meB: continue if len(vals) > 1: pids = [refl_simB[v]['panel'] for v in vals] if len(set(pids)) == 1: print vals # merge_spots(vals) refl_vals = refl_simB.select( flex.bool( [i_v in vals for i_v in np.arange(NreflB)])) x, y, z = spot_utils.xyz_from_refl(refl_vals) allI = [r['intensity.sum.value'] for r in refl_vals] allI = sum(allI) xm = np.mean(x) ym = np.mean(y) zm = np.mean(z) drop_meB.extend(vals[1:]) x1b, x2b, y1b, y2b, z1b, z2b = zip( *[r['bbox'] for r in refl_vals]) keep_me = vals[0] refl_simB['intensity.sum.value'][keep_me] = allI refl_simB['xyzobs.px.value'][keep_me] = (xm, ym, zm) refl_simB['bbox'][keep_me] = (min(x1b), max(x2b), min(y1b),\ max(y2b), min(z1b), max(z2b)) else: drop_meB.append(vals) print vals if drop_meB: keep_meB = [i not in drop_meB for i in range(NreflB)] refl_simB = refl_simB.select(flex.bool(keep_meB)) NreflB = len(refl_simB) ## remake the trees given the drops _, all_HiA = spot_utils.refls_to_hkl(refl_simA, detector, beamA, crystal=crystalAB, returnQ=False) all_treeA = cKDTree(all_HiA) _, all_HiB = spot_utils.refls_to_hkl(refl_simB, detector, beamB, crystal=crystalAB, returnQ=False) #all_treeB = cKDTree(all_HiB) ## CHECK if same HKL, indexed by both colors # exists on multiple panels, and if so, delete... nnAB = all_treeA.query_ball_point(all_HiB, r=1e-7) drop_meA = [] drop_meB = [] for iB, iA_vals in enumerate(nnAB): if len(iA_vals) > 0: assert (len(iA_vals) == 1) iA = iA_vals[0] pidA = refl_simA[iA]['panel'] pidB = refl_simB[iB]['panel'] if pidA != pidB: drop_meA.append(iA) drop_meB.append(iB) if drop_meA: keep_meA = [i not in drop_meA for i in range(NreflA)] refl_simA = refl_simA.select(flex.bool(keep_meA)) if drop_meB: keep_meB = [i not in drop_meB for i in range(NreflB)] refl_simB = refl_simB.select(flex.bool(keep_meB)) # ---- Done with edge case filters# # reflections per panel rpp = spot_utils.refls_by_panelname(refl_data) rppA = spot_utils.refls_by_panelname(refl_simA) rppB = spot_utils.refls_by_panelname(refl_simB) DATA = { "D": [], "IA": [], "IB": [], "h2": [], "k2": [], "l2": [], "h": [], "k": [], "l": [], "PA": [], "PB": [], "FA": [], "FB": [], "iA": [], "iB": [], "Nstrong": [], "pid": [], "delta_pix": [], "deltaX": [], "deltaY": [] } all_int_me = [] # now set up boundboxes and integrate if tilt_plane_integration: mask = np.ones(simsDataSum.shape).astype(np.bool) print "Using tilt plane integration!" else: print "Not using tilt plane integration, just basic spot thresh integration " for pid in rpp: if tilt_plane_integration: Is, Ibk, noise, pix_per = \ integrate.integrate3( rpp[pid], mask[pid], simsDataSum[pid], gain=28) #nom_gain) R = rpp[pid] if pid in rppA: # are there A-channel reflections on this panel inA = True RA = rppA[pid] xA, yA, _ = spot_utils.xyz_from_refl(RA) pointsA = np.array(zip(xA, yA)) HA, HiA, QA = spot_utils.refls_to_hkl(RA, detector, beamA, crystal=crystalAB, returnQ=True) else: inA = False if pid in rppB: # are there B channel reflections on this channel inB = True RB = rppB[pid] xB, yB, _ = spot_utils.xyz_from_refl(RB) pointsB = np.array(zip(xB, yB)) HB, HiB, QB = spot_utils.refls_to_hkl(RB, detector, beamB, crystal=crystalAB, returnQ=True) else: inB = False x, y, _ = spot_utils.xyz_from_refl(R) x = np.array(x) y = np.array(y) panX, panY = detector[pid].get_image_size() mergesA = [] mergesB = [] if inA and inB: # are there both A and B channel reflections ? If so, lets find out which ones have same hkl # make tree structure for merging the spots treeA = cKDTree(pointsA) treeB = cKDTree(pointsB) QA = geom_utils.res_on_panel(detector[pid], beamA) QAmag = np.linalg.norm(QA, axis=2) * 2 * np.pi detdist = detector[pid].get_distance() pixsize = detector[pid].get_pixel_size()[0] merge_me = [] for p in pointsA: iix, iiy = int(p[0]), int(p[1]) q = QAmag[iiy, iix] radA = detdist * np.tan( 2 * np.arcsin(q * waveA / 4 / np.pi)) / pixsize radB = detdist * np.tan( 2 * np.arcsin(q * waveB / 4 / np.pi)) / pixsize rmax = np.abs(radA - radB) split_spot_pairs = treeB.query_ball_point(x=p, r=rmax + sz) merge_me.append(split_spot_pairs) #rmax = geom_utils.twocolor_deltapix(detector[pid], beamA, beamB) #merge_me = treeA.query_ball_tree(treeB, r=rmax + sz) for iA, iB in enumerate(merge_me): if not iB: continue iB = iB[0] # check that the miller indices are the same if not all([i == j for i, j in zip(HiA[iA], HiB[iB])]): continue x1A, x2A, y1A, y2A, _, _ = RA[iA]['bbox'] # shoebox'].bbox x1B, x2B, y1B, y2B, _, _ = RB[iB]['bbox'] # shoebox'].bbox xlow = max([0, min((x1A, x1B)) - sz]) xhigh = min([panX, max((x2A, x2B)) + sz]) ylow = max([0, min((y1A, y1B)) - sz]) yhigh = min([panY, max((y2A, y2B)) + sz]) #if iA==79: # embed() # integrate me if I am in the bounding box! int_me = np.where((xlow < x) & (x < xhigh) & (ylow < y) & (y < yhigh))[0] if not int_me.size: continue mergesA.append(iA) mergesB.append(iB) # integrate the spot, this will change depending on data or simulation totalI = 0 totalCOM = 0 for ref_idx in int_me: if tilt_plane_integration: totalI += Is[ref_idx] else: totalI += rpp[pid][ref_idx]["intensity.sum.value"] totalCOM += np.array( rpp[pid][ref_idx]["xyzobs.px.value"]) totalCOM /= len(int_me) PA = RA[iA]['intensity.sum.value'] PB = RB[iB]['intensity.sum.value'] # get the hkl structure factor, and the sym equiv hkl (h, k, l) = HiA[iA] # NOTE: same for A and B channels (h2, k2, l2), FA = get_val_at_hkl((h, k, l), HA_val_map) _, FB = get_val_at_hkl( (h, k, l), HB_val_map) # NOTE: no need to return h2,k2,l2 twice #if FB==-1 or FA==-1: # continue DATA['h'].append(h) DATA['k'].append(k) DATA['l'].append(l) DATA['h2'].append(h2) DATA['k2'].append(k2) DATA['l2'].append(l2) DATA['D'].append(totalI) DATA['PA'].append(PA) DATA['PB'].append(PB) DATA['FA'].append(FA) DATA['FB'].append(FB) DATA['IA'].append(abs(FA)**2) DATA['IB'].append(abs(FB)**2) DATA['pid'].append(pid) DATA["Nstrong"].append(int_me.size) DATA["iA"].append(iA) DATA["iB"].append(iB) all_int_me.append(int_me) # NOTE: stash the sim-data distance (COM to COM) posA = RA[iA]['xyzobs.px.value'] posB = RB[iB]['xyzobs.px.value'] simCOM = np.mean([posA, posB], axis=0) DATA["delta_pix"].append( distance.euclidean(totalCOM[:2], simCOM[:2])) DATA["deltaX"].append(totalCOM[0] - simCOM[0]) DATA["deltaY"].append(totalCOM[1] - simCOM[1]) if inA: for iA, ref in enumerate(RA): if iA in mergesA: continue x1A, x2A, y1A, y2A, _, _ = RA[iA][ 'bbox'] # ['shoebox'].bbox xlow = max((0, x1A - sz)) xhigh = min((panX, x2A + sz)) ylow = max((0, y1A - sz)) yhigh = min((panY, y2A + sz)) int_me = np.where((xlow < x) & (x < xhigh) & (ylow < y) & (y < yhigh))[0] if not int_me.size: continue totalI = 0 totalCOM = 0 for ref_idx in int_me: if tilt_plane_integration: totalI += Is[ref_idx] else: totalI += rpp[pid][ref_idx]["intensity.sum.value"] totalCOM += np.array( rpp[pid][ref_idx]["xyzobs.px.value"]) totalCOM /= len(int_me) PA = RA[iA]['intensity.sum.value'] PB = 0 # crucial ;) # get the hkl structure factor, and the sym equiv hkl (h, k, l) = HiA[iA] # NOTE: same for A and B channels (h2, k2, l2), FA = get_val_at_hkl((h, k, l), HA_val_map) _, FB = get_val_at_hkl( (h, k, l), HB_val_map) # NOTE: no need to return h2,k2,l2 twice #if FA==-1 or FB==-1: # continue DATA['h'].append(h) DATA['k'].append(k) DATA['l'].append(l) DATA['h2'].append(h2) DATA['k2'].append(k2) DATA['l2'].append(l2) DATA['D'].append(totalI) DATA['PA'].append(PA) DATA['PB'].append(PB) DATA['FA'].append(FA) DATA['FB'].append(FB) DATA['IA'].append(abs(FA)**2) DATA['IB'].append(abs(FB)**2) DATA['pid'].append(pid) DATA["Nstrong"].append(int_me.size) DATA["iA"].append(iA) DATA["iB"].append(np.nan) all_int_me.append(int_me) # NOTE: stash the sim-data distance (COM to COM) simCOM = np.array(RA[iA]['xyzobs.px.value']) DATA["delta_pix"].append( distance.euclidean(totalCOM[:2], simCOM[:2])) DATA["deltaX"].append(totalCOM[0] - simCOM[0]) DATA["deltaY"].append(totalCOM[1] - simCOM[1]) if inB: for iB, ref in enumerate(RB): if iB in mergesB: continue x1B, x2B, y1B, y2B, _, _ = RB[iB]['bbox'] # shoebox'].bbox xlow = max((0, x1B - sz)) xhigh = min((panX, x2B + sz)) ylow = max((0, y1B - sz)) yhigh = min((panY, y2B + sz)) # subimg = simsDataSum[pid][ylow:yhigh, xlow:xhigh] # bg = 0 int_me = np.where((xlow < x) & (x < xhigh) & (ylow < y) & (y < yhigh))[0] if not int_me.size: continue totalI = 0 totalCOM = 0 for ref_idx in int_me: if tilt_plane_integration: totalI += Is[ref_idx] else: totalI += rpp[pid][ref_idx]["intensity.sum.value"] totalCOM += np.array( rpp[pid][ref_idx]["xyzobs.px.value"]) totalCOM /= len(int_me) PA = 0 # crucial ;) PB = RB[iB]['intensity.sum.value'] # get the hkl structure factor, and the sym equiv hkl (h, k, l) = HiB[iB] # NOTE: same for A and B channels (h2, k2, l2), FB = get_val_at_hkl((h, k, l), HB_val_map) _, FA = get_val_at_hkl( (h, k, l), HA_val_map) # NOTE: no need to return h2,k2,l2 twice #if FA==-1 or FB==-1: # continue DATA['h'].append(h) DATA['k'].append(k) DATA['l'].append(l) DATA['h2'].append(h2) DATA['k2'].append(k2) DATA['l2'].append(l2) DATA['D'].append(totalI) DATA['PA'].append(PA) DATA['PB'].append(PB) DATA['FA'].append(FA) DATA['FB'].append(FB) DATA['IA'].append(abs(FA)**2) DATA['IB'].append(abs(FB)**2) DATA['pid'].append(pid) DATA["Nstrong"].append(int_me.size) DATA["iA"].append(np.nan) DATA["iB"].append(iB) all_int_me.append(int_me) # NOTE: stash the sim-data distance (COM to COM) simCOM = np.array(RB[iB]['xyzobs.px.value']) DATA["delta_pix"].append( distance.euclidean(totalCOM[:2], simCOM[:2])) DATA["deltaX"].append(totalCOM[0] - simCOM[0]) DATA["deltaY"].append(totalCOM[1] - simCOM[1]) df = pandas.DataFrame(DATA) df["run"] = run df["shot_idx"] = shot_idx df['gain'] = GAIN if use_data_spec: print "Setting LA, LB as sums over flux regions A,B" df['LA'] = data_fluxes[:75].sum() df['LB'] = data_fluxes[75:].sum() else: print "Setting LA LB as data_fluxes" df['LA'] = data_fluxes[0] df["LB"] = data_fluxes[1] df['K'] = FF[0]**2 * FLUX[0] df["rhs"] = df.gain * (df.IA * df.LA * (df.PA / df.K) + df.IB * df.LB * (df.PB / df.K)) df["lhs"] = df.D #df['data_name'] = data_name df['init_rot'] = init_rot df.to_pickle(pklname) print("PLOT") if args.plot: import pylab as plt plt.plot(df.lhs, df.rhs, '.') plt.show() print("DonDonee")
loader0._h5_handle["space_group_hall_symbol"][()] } cryst_descrB = { '__id__': 'crystal', 'real_space_a': loader1._h5_handle["real_space_a"][()], 'real_space_b': loader1._h5_handle["real_space_b"][()], 'real_space_c': loader1._h5_handle["real_space_c"][()], 'space_group_hall_symbol': loader1._h5_handle["space_group_hall_symbol"][()] } CrystalA = CrystalFactory.from_dict(cryst_descrA) CrystalB = CrystalFactory.from_dict(cryst_descrB) rot_diffs = rotation_matrix_differences((CrystalA, CrystalB)) print(rot_diffs) if args.binary is not None: img0 = img0 > args.binary img1 = img1 > args.binary imgs = cycle([img0, img1]) plt.figure() ax = plt.gca() im = ax.imshow(img0, vmin=args.vmin, vmax=args.vmax, cmap='gnuplot') while 1: im.set_data(imgs.next()) plt.draw()
Cinit = data_orig["crystalAB"] img_f = data_orig["img_f"] loader = dxtbx.load(img_f) cryst_descr = { '__id__': 'crystal', 'real_space_a': loader._h5_handle["real_space_a"][()], 'real_space_b': loader._h5_handle["real_space_b"][()], 'real_space_c': loader._h5_handle["real_space_c"][()], 'space_group_hall_symbol': loader._h5_handle["space_group_hall_symbol"][()] } Csim = CrystalFactory.from_dict(cryst_descr) init_comp = rotation_matrix_differences((Csim, Cinit)) ref_comp = rotation_matrix_differences((Csim, Cmax)) init_rot = float(init_comp.split("\n")[-2].split()[2]) ref_rot = float(ref_comp.split("\n")[-2].split()[2]) print "Orientation deviation changed from %.4f to %.4f during refinement" % ( init_rot, ref_rot) inits.append(init_rot) refs.append(ref_rot) if not args.noabs: inits = np.abs(inits) refs = np.abs(refs) plt.figure() bins = np.linspace(args.brange[0], args.brange[1], args.nbins) plt.hist(inits, bins, histtype='step', lw=2, label="Before refinement")
def main(rank): device_Id = rank % ngpu worker_Id = node_id*ngpu + rank import os import sys from copy import deepcopy import glob from itertools import izip from scipy.spatial import distance import h5py import scipy.ndimage from IPython import embed import numpy as np import pandas from scipy.spatial import cKDTree from simtbx.nanoBragg import shapetype, nanoBragg from libtbx.phil import parse from scitbx.matrix import sqr import dxtbx from dxtbx.model.experiment_list import ExperimentListFactory from dxtbx.model.crystal import CrystalFactory from dials.algorithms.indexing.compare_orientation_matrices \ import rotation_matrix_differences from dials.array_family import flex from dials.command_line.find_spots import phil_scope as find_spots_phil_scope from cxid9114.refine import metrics from cxid9114 import utils from cxid9114.geom import geom_utils from cxid9114.spots import integrate, spot_utils from cxid9114 import parameters from cxid9114.sim import sim_utils from cctbx import miller, sgtbx from cxid9114 import utils from cxid9114.bigsim import sim_spectra from cxid9114.refine.jitter_refine import make_param_list spot_par = find_spots_phil_scope.fetch(source=parse("")).extract() spot_par.spotfinder.threshold.dispersion.global_threshold = 40 spot_par.spotfinder.threshold.dispersion.gain = GAIN spot_par.spotfinder.threshold.dispersion.kernel_size = [2,2] spot_par.spotfinder.threshold.dispersion.sigma_strong = 1 spot_par.spotfinder.threshold.dispersion.sigma_background = 6 spot_par.spotfinder.filter.min_spot_size = 3 spot_par.spotfinder.force_2d = True odir = args.odir odirj = os.path.join(odir, "job%d" % worker_Id) #all_pkl_files = [s for sl in \ # [ files for _,_, files in os.walk(odir)]\ # for s in sl if s.endswith("pkl")] #print "Found %d pkl files already in %s!" \ # % (len(all_pkl_files), odir) if not os.path.exists(odirj): os.makedirs(odirj) hkl_tol = .15 run = 61 shot_idx = 0 ENERGIES = [parameters.ENERGY_LOW, parameters.ENERGY_HIGH] # colors of the beams FF = [10000, None] cryst_descr = {'__id__': 'crystal', 'real_space_a': (79, 0, 0), 'real_space_b': (0, 79, 0), 'real_space_c': (0, 0, 38), 'space_group_hall_symbol': '-P 4 2'} crystalAB = CrystalFactory.from_dict(cryst_descr) sfall_main = sim_spectra.load_spectra("../bigsim/test_sfall.h5") FFdat = [sfall_main[19], sfall_main[110]] FLUX = [1e11, 1e11] # fluxes of the beams chanA_flux = np.random.uniform(1e11,1e12) chanB_flux = np.random.uniform(1e11,1e12) FLUXdat = [chanA_flux, chanB_flux] waveA = parameters.ENERGY_CONV / ENERGIES[0] waveB = parameters.ENERGY_CONV / ENERGIES[1] from cxid9114.bigsim.bigsim_geom import DET,BEAM detector = DET print("Rank %d Begin" % worker_Id) for i_data in range( args.num_trials): pklname = "%s_rank%d_data%d.pkl" % (ofile, worker_Id, i_data) pklname = os.path.join( odirj, pklname) print("<><><><><><><") print("Job %d: trial %d / %d" % ( worker_Id, i_data+1, args.num_trials )) print("<><><><><><><") if (worker_Id==0 and i_data % smi_stride==0 and cuda): print("GPU status") os.system("nvidia-smi") print("\n\n") print("CPU memory usage") mem_usg= """ps -U dermen --no-headers -o rss | awk '{ sum+=$1} END {print int(sum/1024) "MB consumed by CPU user"}'""" os.system(mem_usg) beamA = deepcopy(BEAM) beamB = deepcopy(BEAM) beamA.set_wavelength(waveA) beamB.set_wavelength(waveB) SCALE = np.random.uniform(0.1,10) np.random.seed(args.seed) crystalAB = CrystalFactory.from_dict(cryst_descr) randnums = np.random.random(3) Rrand = random_rotation(1, randnums) crystalAB.set_U(Rrand.ravel()) #pert = np.random.uniform(0.0001/2/np.pi, 0.0003 / 2. /np.pi) #print("PERT %f" % pert) #Rsmall = random_rotation(0.00001, randnums ) #pert) params_lst = make_param_list(crystalAB, DET, BEAM, 1, rot=0.08, cell=.0000001, eq=(1,1,0), min_Ncell=23, max_Ncell=24, min_mos_spread=0.02, max_mos_spread=0.08) Ctruth = params_lst[0]['crystal'] print Ctruth.get_unit_cell().parameters() print crystalAB.get_unit_cell().parameters() init_comp = rotation_matrix_differences((Ctruth, crystalAB)) init_rot = float(init_comp.split("\n")[-2].split()[2]) if use_data_spec: print "NOT IMPLEMENTED, Using a phony 2col spectrum to simulate the data" data_fluxes = FLUXdat data_energies = [parameters.ENERGY_LOW, parameters.ENERGY_HIGH] data_ff = FFdat else: print "Using a phony two color spectrum to simulate the data" data_fluxes = FLUXdat data_energies = [parameters.ENERGY_LOW, parameters.ENERGY_HIGH] data_ff = FFdat print ("Truth crystal Misorientation deviation: %f deg" % init_rot ) if args.truth_cryst: print "Using truth crystal" dataCryst = Ctruth else: print "Not using truth crystal" dataCryst = crystalAB if not make_background: print "SIMULATING Flat-Fhkl IMAGES" simsAB = sim_utils.sim_twocolors2( crystalAB, detector, BEAM, FF, [parameters.ENERGY_LOW, parameters.ENERGY_HIGH], FLUX, pids=None, Gauss=Gauss, cuda=cuda, oversample=oversample, Ncells_abc=Ncells_abc, mos_dom=mos_doms, mos_spread=mos_spread, exposure_s=exposure_s, beamsize_mm=beamsize_mm, device_Id=device_Id, boost=boost) if make_background: print("MAKING BACKGROUND") spec_file = h5py.File("../bigsim/simMe_data_run62.h5", "r") ave_spec = np.mean( spec_file["hist_spec"][()], axis=0) data_fluxes=[ave_spec[19], ave_spec[110] ] data_energies = spec_file["energy_bins"][()][[19,110]] data_ff = [1,1] #*len(data_energies) only_water=True else: only_water=False print "SIULATING DATA IMAGE" print data_fluxes simsDataSum = sim_utils.sim_twocolors2( dataCryst, detector, BEAM, data_ff, data_energies, data_fluxes, pids=None, Gauss=Gauss, cuda=cuda,oversample=oversample, Ncells_abc=Ncells_abc, accumulate=True, mos_dom=mos_doms, mos_spread=mos_spread, boost=boost, exposure_s=exposure_s, beamsize_mm=beamsize_mm, only_water=only_water, device_Id=device_Id) simsDataSum = SCALE*np.array(simsDataSum) if make_background: bg_out = h5py.File(bg_name, "w") bg_out.create_dataset("bigsim_d9114",data=simsDataSum[0]) print "Background made! Saved to file %s" % bg_name sys.exit() if add_background: print("ADDING BG") background = h5py.File(bg_name, "r")['bigsim_d9114'][()] bg_scale = np.sum([39152412349.12075, 32315440627.406036] ) bg_scale = np.sum(data_fluxes) / bg_scale print "%.3e backgorund scale" % bg_scale print "BG shape", background.shape simsDataSum[0] += background * bg_scale if add_noise: print("ADDING NOISE") for pidx in range(1): SIM = nanoBragg(detector=DET, beam=BEAM, panel_id=pidx) SIM.exposure_s = exposure_s SIM.beamsize_mm = beamsize_mm SIM.flux = np.sum(data_fluxes) SIM.detector_psf_kernel_radius_pixels=5; #SIM.detector_psf_type=shapetype.Gauss SIM.detector_psf_type=shapetype.Unknown # for CSPAD SIM.detector_psf_fwhm_mm=0 SIM.quantum_gain = GAIN SIM.raw_pixels = flex.double(simsDataSum[pidx].ravel()) SIM.add_noise() simsDataSum[pidx] = SIM.raw_pixels.as_numpy_array()\ .reshape(simsDataSum[0].shape) SIM.free_all() del SIM if args.write_img: print "SAVING DATAFILE" h5name = "%s_rank%d_data%d.h5" % (ofile, worker_Id, i_data) h5name = os.path.join( odirj, h5name) fout = h5py.File(h5name,"w" ) fout.create_dataset("bigsim_d9114", data=simsDataSum[0]) fout.create_dataset("crystalAB", data=crystalAB.get_A() ) fout.create_dataset("dataCryst", data=dataCryst.get_A() ) fout.close() if args.write_sim_img: print "SAVING DATAFILE" for i_sim in simsAB: sim_h5name = "%s_rank%d_sim%d_%d.h5" % (ofile, worker_Id, i_data, i_sim) sim_h5name = os.path.join( odirj, sim_h5name) fout = h5py.File(sim_h5name,"w" ) fout.create_dataset("bigsim_d9114", data=simsAB[i_sim][0]) fout.create_dataset("crystalAB", data=crystalAB.get_A() ) fout.create_dataset("dataCryst", data=dataCryst.get_A() ) fout.close() print "RELFS FROM SIMS" refl_simA = spot_utils.refls_from_sims(simsAB[0], detector, beamA, thresh=thresh) refl_simB = spot_utils.refls_from_sims(simsAB[1], detector, beamB, thresh=thresh) if use_dials_spotter: print("DIALS SPOTTING") El = utils.explist_from_numpyarrays(simsDataSum,DET,beamA) refl_data = flex.reflection_table.from_observations(El, spot_par) print("Found %d refls using DIALS spot finder" % len(refl_data)) else: refl_data = spot_utils.refls_from_sims(simsDataSum, detector, beamA,\ thresh=thresh) print ("Found %d refls using threshold" % len(refl_data)) if len(refl_data)==0: print "Rank %d: No reflections found! " % (worker_Id) continue residA = metrics.check_indexable2( refl_data, refl_simA, detector, beamA, crystalAB, hkl_tol) residB = metrics.check_indexable2( refl_data, refl_simB, detector, beamB, crystalAB, hkl_tol) FA = sfall_main[19] # utils.open_flex('SA.pkl') # ground truth values FB = sfall_main[110] #utils.open_flex('SB.pkl') # ground truth values HA = tuple([hkl for hkl in FA.indices()]) HB = tuple([hkl for hkl in FB.indices()]) HA_val_map = { h:data for h,data in izip(FA.indices(), FA.data())} HB_val_map = { h:data for h,data in izip(FB.indices(), FB.data())} Hmaps = [HA_val_map, HB_val_map] def get_val_at_hkl(hkl, val_map): sg96 = sgtbx.space_group(" P 4nw 2abw") poss_equivs = [i.h() for i in miller.sym_equiv_indices(sg96, hkl).indices()] in_map=False for hkl2 in poss_equivs: if hkl2 in val_map: # fast lookup in_map=True break if in_map: return hkl2, val_map[hkl2] else: return (None,None,None),-1 if use_data_spec: print "Setting LA, LB as sums over flux regions A,B" LA = data_fluxes[:75].sum() LB = data_fluxes[75:].sum() else: print "Setting LA LB as data_fluxes" LA = data_fluxes[0] LB = data_fluxes[1] K=FF[0] ** 2 * FLUX[0] L_at_color = [LA,LB] out = spot_utils.integrate_boxes( refl_data, simsDataSum[0], [refl_simA, refl_simB], DET, [beamA,beamB], crystalAB, delta_q=args.deltaq, gain=GAIN ) Nh = len(out[0]) rhs = [] lhs = [] all_H2 = [] all_PA = [] all_PB = [] all_FA = [] all_FB = [] for i in range(Nh): HKL = out[0][i] yobs = out[1][i] Pvals = out[2][i] ycalc = 0 for i_P, P in enumerate(Pvals): L = L_at_color[i_P] H2, F = get_val_at_hkl(HKL, Hmaps[i_P]) if i_P==0: all_FA.append(F) else: all_FB.append(F) ycalc += SCALE*L*P*abs(F)**2/K all_PA.append( Pvals[0]) all_PB.append( Pvals[1]) all_H2.append(H2) rhs.append(ycalc) lhs.append(yobs) df = pandas.DataFrame({"rhs":rhs, "lhs": lhs, "PA":all_PA, "PB":all_PB, "FA": all_FA, "FB": all_FB}) df["run"] = run df["shot_idx"] = shot_idx df['gain'] = SCALE df['LA'] = LA df['LB'] = LB df['K'] = K df['init_rot'] = init_rot h,k,l = zip(*all_H2) df['h2'] = h df['k2'] = k df['l2'] = l df.to_pickle(pklname) if args.plot: print("PLOT") import pylab as plt plt.plot( df.lhs, df.rhs, '.') plt.show() print("DonDonee")