def dxtbx_crystal_from_ucell_and_symbol(ucell_tuple_Adeg, symbol): """ :param ucell_tuple_Adeg: unit cell tuple a,b,c al, be, ga in Angstom and degrees :param symbol: lookup symbol for space group, e.g. 'P1' :return:a default crystal in conventional orientation, a along x-axis """ from cctbx import crystal from dxtbx.model.crystal import CrystalFactory symm = crystal.symmetry("%f,%f,%f,%f,%f,%f" % ucell_tuple_Adeg, symbol) ucell = symm.unit_cell() O = ucell.orthogonalization_matrix() real_space_a = O[0], O[3], O[6] real_space_b = O[1], O[4], O[7] real_space_c = O[2], O[5], O[8] hall_symbol = symm.space_group_info().type().hall_symbol() return CrystalFactory.from_dict({ '__id__': 'crystal', 'real_space_a': real_space_a, 'real_space_b': real_space_b, 'real_space_c': real_space_c, 'space_group_hall_symbol': hall_symbol })
def basic_crystal(): print("Make a randomly oriented xtal") # make a randomly oriented crystal.. np.random.seed(3142019) # make random rotation about principle axes x = col((-1, 0, 0)) y = col((0, -1, 0)) z = col((0, 0, -1)) rx, ry, rz = np.random.uniform(-180, 180, 3) RX = x.axis_and_angle_as_r3_rotation_matrix(rx, deg=True) RY = y.axis_and_angle_as_r3_rotation_matrix(ry, deg=True) RZ = z.axis_and_angle_as_r3_rotation_matrix(rz, deg=True) M = RX * RY * RZ real_a = M * col((79, 0, 0)) real_b = M * col((0, 79, 0)) real_c = M * col((0, 0, 38)) # dxtbx crystal description cryst_descr = { '__id__': 'crystal', 'real_space_a': real_a.elems, 'real_space_b': real_b.elems, 'real_space_c': real_c.elems, 'space_group_hall_symbol': ' P 4nw 2abw' } return CrystalFactory.from_dict(cryst_descr)
def crystal(infile): """Load the given JSON file. Params: infile The input filename or file object Returns: The models """ # If the input is a string then open and read from that file if isinstance(infile, str): with open(infile) as infile: return CrystalFactory.from_dict(json.loads(infile.read())) # Otherwise assume the input is a file and read from it else: return CrystalFactory.from_dict(json.loads(infile.read()))
def response_to_xml(d): if "n_spots_total" in d: response = ( """\ <image>%(image)s</image> <spot_count>%(n_spots_total)s</spot_count> <spot_count_no_ice>%(n_spots_no_ice)s</spot_count_no_ice> <d_min>%(estimated_d_min).2f</d_min> <d_min_method_1>%(d_min_distl_method_1).2f</d_min_method_1> <d_min_method_2>%(d_min_distl_method_2).2f</d_min_method_2> <total_intensity>%(total_intensity).0f</total_intensity>""" % d ) else: assert "error" in d return "<response>\n%s\n</response>" % d["error"] if "lattices" in d: from dxtbx.model.crystal import CrystalFactory for lattice in d["lattices"]: crystal = CrystalFactory.from_dict(lattice["crystal"]) response = "\n".join( [ response, "<unit_cell>%.6g %.6g %.6g %.6g %.6g %.6g</unit_cell>" % (crystal.get_unit_cell().parameters()), ] ) response = "\n".join( [ response, "<n_indexed>%i</n_indexed>" % d["n_indexed"], "<fraction_indexed>%.2f</fraction_indexed>" % d["fraction_indexed"], ] ) if "integrated_intensity" in d: response = "\n".join( [ response, "<integrated_intensity>%.0f</integrated_intensity>" % d["integrated_intensity"], ] ) return "<response>\n%s\n</response>" % response
def response_to_xml(d): if "n_spots_total" in d: response = f"""<image>{d['image']}</image> <spot_count>{d['n_spots_total']}</spot_count> <spot_count_no_ice>{d['n_spots_no_ice']}</spot_count_no_ice> <d_min>{d['estimated_d_min']:.2f}</d_min> <d_min_method_1>{d['d_min_distl_method_1']:.2f}</d_min_method_1> <d_min_method_2>{d['d_min_distl_method_2']:.2f}</d_min_method_2> <total_intensity>{d['total_intensity']:.0f}</total_intensity>""" else: assert "error" in d return f"<response>\n{d['error']}\n</response>" if "lattices" in d: for lattice in d["lattices"]: crystal = CrystalFactory.from_dict(lattice["crystal"]) response = "\n".join( [ response, "<unit_cell>%.6g %.6g %.6g %.6g %.6g %.6g</unit_cell>" % (crystal.get_unit_cell().parameters()), ] ) response = "\n".join( [ response, "<n_indexed>%i</n_indexed>" % d["n_indexed"], "<fraction_indexed>%.2f</fraction_indexed>" % d["fraction_indexed"], ] ) if "integrated_intensity" in d: response = "\n".join( [ response, "<integrated_intensity>%.0f</integrated_intensity>" % d["integrated_intensity"], ] ) return f"<response>\n{response}\n</response>"
def response_to_xml(d): if 'n_spots_total' in d: response = '''\ <image>%(image)s</image> <spot_count>%(n_spots_total)s</spot_count> <spot_count_no_ice>%(n_spots_no_ice)s</spot_count_no_ice> <d_min>%(estimated_d_min).2f</d_min> <d_min_method_1>%(d_min_distl_method_1).2f</d_min_method_1> <d_min_method_2>%(d_min_distl_method_2).2f</d_min_method_2> <total_intensity>%(total_intensity).0f</total_intensity>''' % d else: assert 'error' in d return '<response>\n%s\n</response>' % d['error'] if 'lattices' in d: from dxtbx.model.crystal import CrystalFactory for lattice in d['lattices']: crystal = CrystalFactory.from_dict(lattice['crystal']) response = '\n'.join([ response, '<unit_cell>%.6g %.6g %.6g %.6g %.6g %.6g</unit_cell>' % (crystal.get_unit_cell().parameters()) ]) response = '\n'.join([ response, '<n_indexed>%i</n_indexed>' % d['n_indexed'], '<fraction_indexed>%.2f</fraction_indexed>' % d['fraction_indexed'] ]) if 'integrated_intensity' in d: response = '\n'.join([ response, '<integrated_intensity>%.0f</integrated_intensity>' % d['integrated_intensity'] ]) return '<response>\n%s\n</response>' % response
def run(args): from dials.util.options import OptionParser import libtbx.load_env usage = "%s [options] find_spots.expt" % (libtbx.env.dispatcher_name) parser = OptionParser(usage=usage, phil=phil_scope, epilog=help_message) params, options, args = parser.parse_args( show_diff_phil=True, return_unhandled=True ) positions = None if params.positions is not None: with open(params.positions, "rb") as f: positions = flex.vec2_double() for line in f.readlines(): line = ( line.replace("(", " ") .replace(")", "") .replace(",", " ") .strip() .split() ) assert len(line) == 3 i, x, y = [float(l) for l in line] positions.append((x, y)) assert len(args) == 1 json_file = args[0] import json with open(json_file, "rb") as f: results = json.load(f) n_indexed = flex.double() fraction_indexed = flex.double() n_spots = flex.double() n_lattices = flex.double() crystals = [] image_names = flex.std_string() for r in results: n_spots.append(r["n_spots_total"]) image_names.append(str(r["image"])) if "n_indexed" in r: n_indexed.append(r["n_indexed"]) n_lattices.append(len(r["lattices"])) for d in r["lattices"]: from dxtbx.model.crystal import CrystalFactory crystals.append(CrystalFactory.from_dict(d["crystal"])) else: n_indexed.append(0) n_lattices.append(0) if n_indexed.size(): sel = n_spots > 0 fraction_indexed = flex.double(n_indexed.size(), 0) fraction_indexed.set_selected(sel, n_indexed.select(sel) / n_spots.select(sel)) import matplotlib matplotlib.use("Agg") from matplotlib import pyplot red = "#e74c3c" plot = True table = True grid = params.grid from libtbx import group_args from dials.algorithms.spot_finding.per_image_analysis import plot_stats, print_table estimated_d_min = flex.double() d_min_distl_method_1 = flex.double() d_min_distl_method_2 = flex.double() n_spots_total = flex.int() n_spots_no_ice = flex.int() total_intensity = flex.double() for d in results: estimated_d_min.append(d["estimated_d_min"]) d_min_distl_method_1.append(d["d_min_distl_method_1"]) d_min_distl_method_2.append(d["d_min_distl_method_2"]) n_spots_total.append(d["n_spots_total"]) n_spots_no_ice.append(d["n_spots_no_ice"]) total_intensity.append(d["total_intensity"]) stats = group_args( image=image_names, n_spots_total=n_spots_total, n_spots_no_ice=n_spots_no_ice, n_spots_4A=None, n_indexed=n_indexed, fraction_indexed=fraction_indexed, total_intensity=total_intensity, estimated_d_min=estimated_d_min, d_min_distl_method_1=d_min_distl_method_1, d_min_distl_method_2=d_min_distl_method_2, noisiness_method_1=None, noisiness_method_2=None, ) if plot: plot_stats(stats) pyplot.clf() if table: print_table(stats) n_rows = 10 n_rows = min(n_rows, len(n_spots_total)) perm_n_spots_total = flex.sort_permutation(n_spots_total, reverse=True) print("Top %i images sorted by number of spots:" % n_rows) print_table(stats, perm=perm_n_spots_total, n_rows=n_rows) if flex.max(n_indexed) > 0: perm_n_indexed = flex.sort_permutation(n_indexed, reverse=True) print("Top %i images sorted by number of indexed reflections:" % n_rows) print_table(stats, perm=perm_n_indexed, n_rows=n_rows) print("Number of indexed lattices: ", (n_indexed > 0).count(True)) print( "Number with valid d_min but failed indexing: ", ( (d_min_distl_method_1 > 0) & (d_min_distl_method_2 > 0) & (estimated_d_min > 0) & (n_indexed == 0) ).count(True), ) n_bins = 20 spot_count_histogram( n_spots_total, n_bins=n_bins, filename="hist_n_spots_total.png", log=True ) spot_count_histogram( n_spots_no_ice, n_bins=n_bins, filename="hist_n_spots_no_ice.png", log=True ) spot_count_histogram( n_indexed.select(n_indexed > 0), n_bins=n_bins, filename="hist_n_indexed.png", log=False, ) if len(crystals): plot_unit_cell_histograms(crystals) if params.stereographic_projections and len(crystals): from dxtbx.model.experiment_list import ExperimentListFactory experiments = ExperimentListFactory.from_filenames( [image_names[0]], verbose=False ) assert len(experiments) == 1 imageset = experiments.imagesets()[0] s0 = imageset.get_beam().get_s0() # XXX what if no goniometer? rotation_axis = imageset.get_goniometer().get_rotation_axis() indices = ((1, 0, 0), (0, 1, 0), (0, 0, 1)) for i, index in enumerate(indices): from cctbx import crystal, miller from scitbx import matrix miller_indices = flex.miller_index([index]) symmetry = crystal.symmetry( unit_cell=crystals[0].get_unit_cell(), space_group=crystals[0].get_space_group(), ) miller_set = miller.set(symmetry, miller_indices) d_spacings = miller_set.d_spacings() d_spacings = d_spacings.as_non_anomalous_array().expand_to_p1() d_spacings = d_spacings.generate_bijvoet_mates() miller_indices = d_spacings.indices() # plane normal d0 = matrix.col(s0).normalize() d1 = d0.cross(matrix.col(rotation_axis)).normalize() d2 = d1.cross(d0).normalize() reference_poles = (d0, d1, d2) from dials.command_line.stereographic_projection import ( stereographic_projection, ) projections = [] for cryst in crystals: reciprocal_space_points = ( list(cryst.get_U() * cryst.get_B()) * miller_indices.as_vec3_double() ) projections.append( stereographic_projection(reciprocal_space_points, reference_poles) ) # from dials.algorithms.indexing.compare_orientation_matrices import \ # difference_rotation_matrix_and_euler_angles # R_ij, euler_angles, cb_op = difference_rotation_matrix_and_euler_angles( # crystals[0], cryst) # print max(euler_angles) from dials.command_line.stereographic_projection import plot_projections plot_projections(projections, filename="projections_%s.png" % ("hkl"[i])) pyplot.clf() def plot_grid( values, grid, file_name, cmap=pyplot.cm.Reds, vmin=None, vmax=None, invalid="white", ): values = values.as_double() # At DLS, fast direction appears to be largest direction if grid[0] > grid[1]: values.reshape(flex.grid(reversed(grid))) values = values.matrix_transpose() else: values.reshape(flex.grid(grid)) f, ax1 = pyplot.subplots(1) mesh1 = ax1.pcolormesh(values.as_numpy_array(), cmap=cmap, vmin=vmin, vmax=vmax) mesh1.cmap.set_under(color=invalid, alpha=None) mesh1.cmap.set_over(color=invalid, alpha=None) ax1.set_aspect("equal") ax1.invert_yaxis() pyplot.colorbar(mesh1, ax=ax1) pyplot.savefig(file_name, dpi=600) pyplot.clf() def plot_positions( values, positions, file_name, cmap=pyplot.cm.Reds, vmin=None, vmax=None, invalid="white", ): values = values.as_double() assert positions.size() >= values.size() positions = positions[: values.size()] if vmin is None: vmin = flex.min(values) if vmax is None: vmax = flex.max(values) x, y = positions.parts() dx = flex.abs(x[1:] - x[:-1]) dy = flex.abs(y[1:] - y[:-1]) dx = dx.select(dx > 0) dy = dy.select(dy > 0) scale = 1 / flex.min(dx) # print scale x = (x * scale).iround() y = (y * scale).iround() from libtbx.math_utils import iceil z = flex.double(flex.grid(iceil(flex.max(y)) + 1, iceil(flex.max(x)) + 1), -2) # print z.all() for x_, y_, z_ in zip(x, y, values): z[y_, x_] = z_ plot_grid( z.as_1d(), z.all(), file_name, cmap=cmap, vmin=vmin, vmax=vmax, invalid=invalid, ) return if grid is not None or positions is not None: if grid is not None: positions = tuple(reversed(grid)) plotter = plot_grid else: plotter = plot_positions cmap = pyplot.get_cmap(params.cmap) plotter( n_spots_total, positions, "grid_spot_count_total.png", cmap=cmap, invalid=params.invalid, ) plotter( n_spots_no_ice, positions, "grid_spot_count_no_ice.png", cmap=cmap, invalid=params.invalid, ) plotter( total_intensity, positions, "grid_total_intensity.png", cmap=cmap, invalid=params.invalid, ) if flex.max(n_indexed) > 0: plotter( n_indexed, positions, "grid_n_indexed.png", cmap=cmap, invalid=params.invalid, ) plotter( fraction_indexed, positions, "grid_fraction_indexed.png", cmap=cmap, vmin=0, vmax=1, invalid=params.invalid, ) for i, d_min in enumerate( (estimated_d_min, d_min_distl_method_1, d_min_distl_method_2) ): vmin = flex.min(d_min.select(d_min > 0)) vmax = flex.max(d_min) cmap = pyplot.get_cmap("%s_r" % params.cmap) d_min.set_selected(d_min <= 0, vmax) if i == 0: plotter( d_min, positions, "grid_d_min.png", cmap=cmap, vmin=vmin, vmax=vmax, invalid=params.invalid, ) else: plotter( d_min, positions, "grid_d_min_method_%i.png" % i, cmap=cmap, vmin=vmin, vmax=vmax, invalid=params.invalid, ) if flex.max(n_indexed) > 0: pyplot.hexbin(n_spots, n_indexed, bins="log", cmap=pyplot.cm.jet, gridsize=50) pyplot.colorbar() xlim = pyplot.xlim() ylim = pyplot.ylim() pyplot.plot([0, max(n_spots)], [0, max(n_spots)], c=red) pyplot.xlim(0, xlim[1]) pyplot.ylim(0, ylim[1]) pyplot.xlabel("# spots") pyplot.ylabel("# indexed") pyplot.savefig("n_spots_vs_n_indexed.png") pyplot.clf() pyplot.hexbin( n_spots, fraction_indexed, bins="log", cmap=pyplot.cm.jet, gridsize=50 ) pyplot.colorbar() pyplot.xlim(0, pyplot.xlim()[1]) pyplot.ylim(0, pyplot.ylim()[1]) pyplot.xlabel("# spots") pyplot.ylabel("Fraction indexed") pyplot.savefig("n_spots_vs_fraction_indexed.png") pyplot.clf() pyplot.hexbin( n_indexed, fraction_indexed, bins="log", cmap=pyplot.cm.jet, gridsize=50 ) pyplot.colorbar() pyplot.xlim(0, pyplot.xlim()[1]) pyplot.ylim(0, pyplot.ylim()[1]) pyplot.xlabel("# indexed") pyplot.ylabel("Fraction indexed") pyplot.savefig("n_indexed_vs_fraction_indexed.png") pyplot.clf() pyplot.hexbin(n_spots, n_lattices, bins="log", cmap=pyplot.cm.jet, gridsize=50) pyplot.colorbar() pyplot.xlim(0, pyplot.xlim()[1]) pyplot.ylim(0, pyplot.ylim()[1]) pyplot.xlabel("# spots") pyplot.ylabel("# lattices") pyplot.savefig("n_spots_vs_n_lattices.png") pyplot.clf() pyplot.hexbin( estimated_d_min, d_min_distl_method_1, bins="log", cmap=pyplot.cm.jet, gridsize=50, ) pyplot.colorbar() # pyplot.gca().set_aspect('equal') xlim = pyplot.xlim() ylim = pyplot.ylim() m = max(max(estimated_d_min), max(d_min_distl_method_1)) pyplot.plot([0, m], [0, m], c=red) pyplot.xlim(0, xlim[1]) pyplot.ylim(0, ylim[1]) pyplot.xlabel("estimated_d_min") pyplot.ylabel("d_min_distl_method_1") pyplot.savefig("d_min_vs_distl_method_1.png") pyplot.clf() pyplot.hexbin( estimated_d_min, d_min_distl_method_2, bins="log", cmap=pyplot.cm.jet, gridsize=50, ) pyplot.colorbar() # pyplot.gca().set_aspect('equal') xlim = pyplot.xlim() ylim = pyplot.ylim() m = max(max(estimated_d_min), max(d_min_distl_method_2)) pyplot.plot([0, m], [0, m], c=red) pyplot.xlim(0, xlim[1]) pyplot.ylim(0, ylim[1]) pyplot.xlabel("estimated_d_min") pyplot.ylabel("d_min_distl_method_2") pyplot.savefig("d_min_vs_distl_method_2.png") pyplot.clf() pyplot.hexbin( d_min_distl_method_1, d_min_distl_method_2, bins="log", cmap=pyplot.cm.jet, gridsize=50, ) pyplot.colorbar() # pyplot.gca().set_aspect('equal') xlim = pyplot.xlim() ylim = pyplot.ylim() m = max(max(d_min_distl_method_1), max(d_min_distl_method_2)) pyplot.plot([0, m], [0, m], c=red) pyplot.xlim(0, xlim[1]) pyplot.ylim(0, ylim[1]) pyplot.xlabel("d_min_distl_method_1") pyplot.ylabel("d_min_distl_method_2") pyplot.savefig("distl_method_1_vs_distl_method_2.png") pyplot.clf() pyplot.hexbin(n_spots, estimated_d_min, bins="log", cmap=pyplot.cm.jet, gridsize=50) pyplot.colorbar() pyplot.xlim(0, pyplot.xlim()[1]) pyplot.ylim(0, pyplot.ylim()[1]) pyplot.xlabel("# spots") pyplot.ylabel("estimated_d_min") pyplot.savefig("n_spots_vs_d_min.png") pyplot.clf() pyplot.hexbin( n_spots, d_min_distl_method_1, bins="log", cmap=pyplot.cm.jet, gridsize=50 ) pyplot.colorbar() pyplot.xlim(0, pyplot.xlim()[1]) pyplot.ylim(0, pyplot.ylim()[1]) pyplot.xlabel("# spots") pyplot.ylabel("d_min_distl_method_1") pyplot.savefig("n_spots_vs_distl_method_1.png") pyplot.clf() pyplot.hexbin( n_spots, d_min_distl_method_2, bins="log", cmap=pyplot.cm.jet, gridsize=50 ) pyplot.colorbar() pyplot.xlim(0, pyplot.xlim()[1]) pyplot.ylim(0, pyplot.ylim()[1]) pyplot.xlabel("# spots") pyplot.ylabel("d_min_distl_method_2") pyplot.savefig("n_spots_vs_distl_method_2.png") pyplot.clf()
def main(): from cxid9114.sim import sim_utils from dxtbx.model.crystal import CrystalFactory from dxtbx_model_ext import flex_Beam from dxtbx.model.detector import DetectorFactory from dxtbx.model.beam import BeamFactory from simtbx.nanoBragg.tst_nanoBragg_basic import fcalc_from_pdb import numpy as np from cxid9114.parameters import ENERGY_CONV energies = np.arange(8920, 8930) fluxes = np.ones(len(energies)) * 5e11 patt_args = {"Ncells_abc": (20, 20, 20), "profile": "square", "verbose": 0} beam_descr = { 'direction': (0.0, 0.0, 1.0), 'divergence': 0.0, 'flux': 5e11, 'polarization_fraction': 1., 'polarization_normal': (0.0, 1.0, 0.0), 'sigma_divergence': 0.0, 'transmission': 1.0, 'wavelength': ENERGY_CONV / energies[0] } 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' } det_descr = { 'panels': [{ 'fast_axis': (-1.0, 0.0, 0.0), 'gain': 1.0, 'identifier': '', 'image_size': (196, 196), 'mask': [], 'material': '', 'mu': 0.0, 'name': 'Panel', 'origin': (19.6, -19.6, -550), 'pedestal': 0.0, 'pixel_size': (0.1, 0.1), 'px_mm_strategy': { 'type': 'SimplePxMmStrategy' }, 'raw_image_offset': (0, 0), 'slow_axis': (0.0, 1.0, 0.0), 'thickness': 0.0, 'trusted_range': (0.0, 65536.0), 'type': '' }] } DET = DetectorFactory.from_dict(det_descr) BEAM = BeamFactory.from_dict(beam_descr) crystal = CrystalFactory.from_dict(cryst_descr) Patt = sim_utils.PatternFactory(crystal=crystal, detector=DET, beam=BEAM, **patt_args) img = None Fens = [] xrbeams = flex_Beam() for fl, en in zip(fluxes, energies): wave = ENERGY_CONV / en F = fcalc_from_pdb(resolution=4, algorithm="fft", wavelength=wave) Patt.primer(crystal, energy=en, flux=fl, F=F) if img is None: img = Patt.sim_rois(reset=True) # defaults to full detector else: img += Patt.sim_rois(reset=True) Fens.append(F) xrb = BeamFactory.from_dict(beam_descr) xrb.set_wavelength(wave * 1e-10) # need to fix the necessity to do this.. xrb.set_flux(fl) xrb.set_direction(BEAM.get_direction()) xrbeams.append(xrb) #import pylab as plt #def plot_img(ax,img): # m = img[img >0].mean() # s = img[img > 0].std() # vmax = m+5*s # vmin = 0 # ax.imshow(img, vmin=vmin, vmax=vmax, cmap='gnuplot') print("\n\n\n") print("<><><><><><><><><><>") print("NEXT TRIAL") print("<><><><><><><><><><>") # patt2 = sim_utils.PatternFactory(crystal=crystal, detector=DET, beam=xrbeams, **patt_args) patt2.prime_multi_Fhkl(multisource_Fhkl=Fens) img2 = patt2.sim_rois(reset=True) #plt.figure() #ax1 = plt.gca() #plt.figure() #ax2 = plt.gca() #plot_img(ax1, img) #plot_img(ax2, img2) #plt.show() assert (np.allclose(img, img2))
'material': '', 'mu': 0.0, 'name': 'Panel', 'origin': (-im_shape[0]*pixsize/2., im_shape[1]*pixsize/2., -detdist), 'pedestal': 0.0, 'pixel_size': (pixsize, pixsize), 'px_mm_strategy': {'type': 'SimplePxMmStrategy'}, 'raw_image_offset': (0, 0), 'thickness': 0.0, 'trusted_range': (-1e7, 1e7), 'type': ''}]} # make the dxtbx objects BEAM = BeamFactory.from_dict(beam_descr) DETECTOR = DetectorFactory.from_dict(det_descr) CRYSTAL = CrystalFactory.from_dict(cryst_descr) # make a dummie HKL table with constant HKL intensity # this is just to make spots DEFAULT_F = 1e2 symbol = CRYSTAL.get_space_group().info().type().lookup_symbol() # this is just P43212 sgi = sgtbx.space_group_info(symbol) symm = symmetry(unit_cell=CRYSTAL.get_unit_cell(), space_group_info=sgi) miller_set = symm.build_miller_set(anomalous_flag=True, d_min=1.6, d_max=999) Famp = flex.double(np.ones(len(miller_set.indices())) * DEFAULT_F) Famp = miller.array(miller_set=miller_set, data=Famp).set_observation_type_xray_amplitude() Ncells_abc = 20, 20, 20 oversmaple = 2 # do the simulation
S.default_F = 100 S.F000 = 0 S.xtal_shape = simtbx.nanoBragg.shapetype.Tophat # makes nice spots with no shape transforms S.oversample=2 S.Ncells_abc = (10,10,10) S.add_nanoBragg_spots() return [S.raw_pixels.as_numpy_array() ] # Here is our dxtbx crystal description # this will be our ground truth cryst_descr = {'__id__': 'crystal', 'real_space_a': (150., 0, 0), 'real_space_b': (0, 200., 0), 'real_space_c': (0, 0, 100.), 'space_group_hall_symbol': '-P 2 2'} cryst = CrystalFactory.from_dict(cryst_descr) # Lets define the indexing parameters for when we need them later: idxpar = indexer_phil_scope.extract() idxpar.indexing.known_symmetry.space_group = cryst.get_space_group().info() idxpar.indexing.known_symmetry.unit_cell = cryst.get_unit_cell() idxpar.indexing.method = "fft1d" idxpar.indexing.fft1d.characteristic_grid = 0.029 idxpar.indexing.multiple_lattice_search.max_lattices = 1 idxpar.indexing.stills.indexer = 'stills' idxpar.indexing.stills.refine_all_candidates = True idxpar.indexing.stills.refine_candidates_with_known_symmetry = True idxpar.indexing.stills.candidate_outlier_rejection = False idxpar.indexing.debug = False idxpar.refinement.verbosity = 0
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")
def __init__(self, params): import cPickle as pickle from dxtbx.model import BeamFactory from dxtbx.model import DetectorFactory from dxtbx.model.crystal import CrystalFactory from cctbx.crystal_orientation import crystal_orientation, basis_type from dxtbx.model import Experiment, ExperimentList from scitbx import matrix self.experiments = ExperimentList() self.unique_file_names = [] self.params = params data = pickle.load( open(self.params.output.prefix + "_frame.pickle", "rb")) frames_text = data.split("\n") for item in frames_text: tokens = item.split(' ') wavelength = float(tokens[order_dict["wavelength"]]) beam = BeamFactory.simple(wavelength=wavelength) detector = DetectorFactory.simple( sensor=DetectorFactory.sensor( "PAD"), # XXX shouldn't hard code for XFEL distance=float(tokens[order_dict["distance"]]), beam_centre=[ float(tokens[order_dict["beam_x"]]), float(tokens[order_dict["beam_y"]]) ], fast_direction="+x", slow_direction="+y", pixel_size=[self.params.pixel_size, self.params.pixel_size], image_size=[1795, 1795], # XXX obviously need to figure this out ) reciprocal_matrix = matrix.sqr([ float(tokens[order_dict[k]]) for k in [ 'res_ori_1', 'res_ori_2', 'res_ori_3', 'res_ori_4', 'res_ori_5', 'res_ori_6', 'res_ori_7', 'res_ori_8', 'res_ori_9' ] ]) ORI = crystal_orientation(reciprocal_matrix, basis_type.reciprocal) direct = matrix.sqr(ORI.direct_matrix()) transfer_dict = dict( __id__="crystal", ML_half_mosaicity_deg=float( tokens[order_dict["half_mosaicity_deg"]]), ML_domain_size_ang=float( tokens[order_dict["domain_size_ang"]]), real_space_a=matrix.row(direct[0:3]), real_space_b=matrix.row(direct[3:6]), real_space_c=matrix.row(direct[6:9]), space_group_hall_symbol=self.params.target_space_group.type(). hall_symbol(), ) crystal = CrystalFactory.from_dict(transfer_dict) """ old code reflects python-based crystal model crystal = Crystal( real_space_a = matrix.row(direct[0:3]), real_space_b = matrix.row(direct[3:6]), real_space_c = matrix.row(direct[6:9]), space_group_symbol = self.params.target_space_group.type().lookup_symbol(), mosaicity = float(tokens[order_dict["half_mosaicity_deg"]]), ) crystal.domain_size = float(tokens[order_dict["domain_size_ang"]]) """ #if isoform is not None: # newB = matrix.sqr(isoform.fractionalization_matrix()).transpose() # crystal.set_B(newB) self.experiments.append( Experiment( beam=beam, detector=None, #dummy for now crystal=crystal)) self.unique_file_names.append( tokens[order_dict["unique_file_name"]]) self.show_summary()
'real_space_b': loader0._h5_handle["real_space_b"][()], 'real_space_c': loader0._h5_handle["real_space_c"][()], 'space_group_hall_symbol': 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')
def main(rank): device_Id = rank % ngpu_per_node worker_Id = rank #node_id*ngpu_per_node + rank # support for running on multiple N-node GPUs to increase worker pool import os import sys import h5py import numpy as np from simtbx.nanoBragg import shapetype, nanoBragg from scitbx.array_family import flex from dxtbx.model.crystal import CrystalFactory from dials.algorithms.indexing.compare_orientation_matrices \ import rotation_matrix_differences, difference_rotation_matrix_axis_angle from cxid9114.sf import struct_fact_special from cxid9114 import parameters from cxid9114.utils import random_rotation from cxid9114.sim import sim_utils #from cxid9114.refine.jitter_refine import make_param_list from cxid9114.geom.single_panel import DET, BEAM if args.cspad: # use a cspad for simulation from cxid9114.geom.multi_panel import CSPAD #from dxtbx.model import Panel # put this new CSPAD in the same plane as the single panel detector (originZ) for pid in range(len(CSPAD)): CSPAD[pid].set_trusted_range(DET[0].get_trusted_range()) #node_d = CSPAD[pid].to_dict() #node_d["origin"] = node_d["origin"][0], node_d["origin"][1], DET[0].get_origin()[2] #CSPAD[pid] = Panel.from_dict(node_d) #from IPython import embed #embed() DET = CSPAD # rename if args.use_rank_as_seed: np.random.seed(worker_Id) else: np.random.seed(args.seed) sim_path = os.path.dirname(sim_utils.__file__) spectra_file = os.path.join(sim_path, "../spec/realspec.h5") sfall_file = os.path.join(sim_path, "../sf/realspec_sfall.h5") data_fluxes_all = h5py.File(spectra_file, "r")["hist_spec"][()] / exposure_s ave_flux_across_exp = np.mean(data_fluxes_all, axis=0).sum() data_sf, data_energies = struct_fact_special.load_sfall(sfall_file) if args.datasf: print("Loading 4bs7 structure factors!") #data_sf = struct_fact_special.load_4bs7_sf() data_sf = struct_fact_special.load_p9() data_sf = [data_sf] + [None] * (len(data_energies) - 1) sad_idx = np.argmin(np.abs(data_energies - 8944)) sad_range = range(sad_idx - 10, sad_idx + 10) n_sad = len(sad_range) if args.sad: data_sf = [data_sf[0]] + [None] * (n_sad - 1) data_energies = np.array([data_energies[i] for i in sad_range]) beamsize_mm = np.sqrt(np.pi * (beam_diam_mm / 2)**2) # TODO: verify this works for all variants of parallelization (multi 8-GPU nodes, multi kernels per GPU etc) data_fluxes_worker = np.array_split( data_fluxes_all, ngpu_per_node * kernels_per_gpu * num_nodes)[worker_Id] a, b, c, _, _, _ = data_sf[0].unit_cell().parameters() hall = data_sf[0].space_group_info().type().hall_symbol() # Each rank (worker) gets its own output directory odir = args.odir odirj = os.path.join(odir, "job%d" % worker_Id) if not os.path.exists(odirj): os.makedirs(odirj) print("Rank %d Begin" % worker_Id) for i_data in range(args.num_trials): h5name = "%s_rank%d_data%d.h5" % (ofile, worker_Id, i_data) h5name = os.path.join(odirj, h5name) if os.path.exists(h5name) and not args.overwrite: print("Job %d: skipping- image %s already exists!" \ %(worker_Id, h5name)) continue print("<><><><><><><") print("Job %d: trial %d / %d" % (worker_Id, i_data + 1, args.num_trials)) print("<><><><><><><") # If im the zeroth worker I want to show usage statisitcs: 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) data_fluxes = data_fluxes_worker[i_data] if args.sad: flux_sum = data_fluxes.sum() data_fluxes = data_fluxes[sad_range] data_fluxes = data_fluxes / data_fluxes.sum() * flux_sum a = np.random.normal(a, args.ucelljitter) c = np.random.normal(c, args.ucelljitter) cryst_descr = { '__id__': 'crystal', # its al,be,ga = 90,90,90 'real_space_a': (a, 0, 0), 'real_space_b': (0, a, 0), 'real_space_c': (0, 0, c), 'space_group_hall_symbol': hall } crystal = CrystalFactory.from_dict(cryst_descr) # generate a random scale factor within two orders of magnitude to apply to the image data.. # rotate the crystal using a known rotation #crystal = CrystalFactory.from_dict(cryst_descr) randnums = np.random.random(3) Rrand = random_rotation(1, randnums) crystal.set_U(Rrand.ravel()) # NOTE: This can be used to simulate jitter in the # crystal model, but for now we disable jitter #params_lst = make_param_list(crystal, 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'] Ncells = int(np.random.normal(args.Ncells, args.Ncellsjitter)) Ncells_abc = (Ncells, Ncells, Ncells) # the following will override some parameters # to aid simulation of a background image using same pipeline if make_background: print("MAKING BACKGROUND : just at two colors") data_fluxes = [ave_flux_across_exp * .5, ave_flux_across_exp * .5] data_energies = [parameters.ENERGY_LOW, parameters.ENERGY_HIGH] # should be 8944 and 9034 data_sf = [ 1, 1 ] # dont care about structure factors when simulating background water scatter Ncells_abc = 10, 10, 10 only_water = True else: only_water = False xtal_size_mm = args.xtal_size_mm if args.xtal_size_jitter is not None: xtal_size_mm = np.random.normal(xtal_size_mm, args.xtal_size_jitter) sim_args = [crystal, DET, BEAM, data_sf, data_energies, data_fluxes] masterscale = args.masterscale if masterscale is not None: masterscale = np.random.normal(masterscale, args.masterscalejitter) if masterscale < 0.1: # FIXME: make me more smart master_scale = 0.1 sim_kwargs = { 'pids': None, 'profile': args.profile, 'cuda': cuda, 'oversample': args.oversample, 'Ncells_abc': Ncells_abc, 'accumulate': True, 'mos_dom': mos_doms, 'mos_spread': mos_spread, 'exposure_s': exposure_s, 'beamsize_mm': beamsize_mm, 'only_water': only_water, 'device_Id': device_Id, 'div_tup': div_tup, 'master_scale': masterscale, 'gimmie_Patt': True, 'adc_offset': adc_offset, 'show_params': args.show_params, 'crystal_size_mm': xtal_size_mm, 'one_sf_array': data_sf[0] is not None and data_sf[1] is None } print("SIULATING DATA IMAGE") if args.optimize_oversample: oversample = 1 reference = None while 1: sim_kwargs['oversample'] = oversample simsDataSum, PattFact = sim_utils.sim_colors( *sim_args, **sim_kwargs) simsDataSum = np.array(simsDataSum) if oversample == 1: reference = simsDataSum else: residual = np.abs(reference - simsDataSum) where_one_photon = np.where(simsDataSum > 1) N_over_sigma = (residual[where_one_photon] > np.sqrt( reference[where_one_photon])).sum() print ("Oversample = %d; |Residual| summed = %f; N over sigma: %d" \ % (oversample, residual.sum(), N_over_sigma)) if np.allclose(simsDataSum, reference, atol=4): print( "Optimal oversample for current parameters: %d\nGoodbyw" % oversample) sys.exit() reference = simsDataSum oversample += 1 else: simsDataSum, PattFact = sim_utils.sim_colors( *sim_args, **sim_kwargs) # spot_scale = PattFact.spot_scale simsDataSum = np.array(simsDataSum) if args.cspad: assert simsDataSum.shape == (64, 185, 194) if make_background: if args.cspad: bg_out = h5py.File(bg_name, "w") bg_out.create_dataset("bigsim_d9114", data=simsDataSum) np.savez("testbg.npz", img=simsDataSum, det=DET.to_dict(), beam=BEAM.to_dict()) else: bg_out = h5py.File(bg_name, "w") bg_out.create_dataset("bigsim_d9114", data=simsDataSum[0]) np.savez("testbg_mono.npz", img=simsDataSum[0], det=DET.to_dict(), beam=BEAM.to_dict()) print("Background made! Saved to file %s" % bg_name) # force an exit here if making a background... sys.exit() if add_background: print("ADDING BG") background = h5py.File(bg_name, "r")['bigsim_d9114'][()] if args.cspad: assert background.shape == (64, 185, 194) # background was made using average flux over all shots, so scale it up/down here bg_scale = data_fluxes.sum() / ave_flux_across_exp # TODO consider varying the background level to simultate jet thickness jitter if args.cspad: simsDataSum += background * bg_scale else: simsDataSum[0] += background * bg_scale if add_noise: print("ADDING NOISE") for pidx in range(len(DET)): SIM = nanoBragg(detector=DET, beam=BEAM, panel_id=pidx) SIM.beamsize_mm = beamsize_mm SIM.exposure_s = exposure_s SIM.flux = np.sum(data_fluxes) SIM.adc_offset_adu = adc_offset #SIM.detector_psf_kernel_radius_pixels = 5 #SIM.detector_psf_type = shapetype.Unknown # for CSPAD SIM.detector_psf_fwhm_mm = 0 SIM.quantum_gain = GAIN if not args.readoutnoise: SIM.readout_noise_adu = 0 SIM.raw_pixels = flex.double(simsDataSum[pidx].ravel()) SIM.add_noise() simsDataSum[pidx] = SIM.raw_pixels.as_numpy_array()\ .reshape(simsDataSum[pidx].shape) SIM.free_all() del SIM #all_rots =[] #from scitbx.matrix import sqr #for i_U, U in enumerate(Umats): # Utruth = crystal.get_U() # crystal2 = CrystalFactory.from_dict(cryst_descr) # crystal2.set_U( sqr(U)*sqr(Utruth) ) # out = difference_rotation_matrix_axis_angle(crystal, crystal2) # all_rots.append(out[2]) #assert( np.mean(all_rots) < 1e-7) if args.write_img: print "SAVING DATAFILE" if args.cspad: np.savez(h5name + ".npz", img=simsDataSum.astype(np.float32), det=DET.to_dict(), beam=BEAM.to_dict()) else: np.savez(h5name + ".npz", img=simsDataSum[0].astype(np.float32), det=DET.to_dict(), beam=BEAM.to_dict()) fout = h5py.File(h5name, "w") #fout.create_dataset("bigsim_d9114", data=simsDataSum[0].astype(np.float32), compression='lzf') fout.create_dataset("crystalA", data=crystal.get_A()) fout.create_dataset("crystalU", data=crystal.get_U()) fout.create_dataset("spectrum", data=data_fluxes) fout.create_dataset("mos_doms", data=mos_doms) fout.create_dataset("mos_spread", data=mos_spread) fout.create_dataset("Ncells_abc", data=Ncells_abc) fout.create_dataset("divergence_tuple", data=div_tup) fout.create_dataset("beamsize_mm", data=beamsize_mm) fout.create_dataset("exposure_s", data=exposure_s) fout.create_dataset("profile", data=profile) fout.create_dataset("xtal_size_mm", data=xtal_size_mm) fout.create_dataset("spot_scale", data=spot_scale) fout.create_dataset("gain", data=GAIN) fout.close() print("DonDonee")
def run_sim2smv(Nshot_max, odir, tag, rank, n_jobs, save_bragg=False, save_smv=True, save_h5 =False, return_pixels=False): import os import h5py import math import sys import numpy as np from IPython import embed from cxid9114.bigsim.bigsim_geom import DET,BEAM import scitbx from scitbx.array_family import flex from scitbx.matrix import sqr,col from simtbx.nanoBragg import shapetype from simtbx.nanoBragg import nanoBragg import libtbx.load_env # possibly implicit from libtbx.development.timers import Profiler from cctbx import crystal,crystal_orientation from LS49.sim.step4_pad import microcrystal from cxid9114 import utils from cxid9114.sim import sim_utils from cxid9114.parameters import ENERGY_CONV, ENERGY_HIGH, ENERGY_LOW from cxid9114.bigsim import sim_spectra from dxtbx.model.crystal import CrystalFactory odir_j = os.path.join( odir, "job%d" % rank) if not os.path.exists(odir_j): os.makedirs(odir_j) 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'} Crystal = CrystalFactory.from_dict(cryst_descr) mos_spread_deg=0.015 mos_doms=1000 beam_size_mm=0.001 exposure_s=1 use_microcrystal=True Deff_A = 2200 length_um = Deff_A/1000. timelog = False if add_background: background = utils.open_flex("background").as_numpy_array() crystal = microcrystal(Deff_A = Deff_A, length_um = length_um, beam_diameter_um = beam_size_mm*1000, verbose=False) spec_file = h5py.File("simMe_data_run62.h5", "r") spec_data = spec_file["hist_spec"] Umat_data = spec_file["Umats"] en_chans = spec_file["energy_bins"][()] ilow = np.abs(en_chans - ENERGY_LOW).argmin() ihigh = np.abs(en_chans - ENERGY_HIGH).argmin() wave_chans = ENERGY_CONV/en_chans sfall_main = sim_spectra.load_spectra("test_sfall.h5") Nshot = spec_data.shape[0] idx_range = np.array_split(np.arange(Nshot), n_jobs) Nshot_per_job = len(idx_range[rank]) if Nshot_max > 0 : Nshot_per_job = min( Nshot_max, Nshot_per_job) print ("Job %d: Simulating %d shots" % (rank, Nshot_per_job)) istart = idx_range[rank][0] istop = istart + Nshot_per_job smi_stride = 10 for idx in range( istart, istop): print ("<><><><><><><><><><><><><><>") print ("Job %d; Image %d (%d - %d)" % (rank, idx+1, istart, istop)) print ("<><><><><><><><><><><><><><>") smv_fileout = os.path.join( odir_j, "%s_%d.img" % (tag,idx)) h5_fileout = smv_fileout + ".h5" if os.path.exists(smv_fileout) and not overwrite and save_smv: print("Shot %s exists: moving on" % smv_fileout) continue if os.path.exists(h5_fileout) and not overwrite and save_h5: print("Shot %s exists: moving on" % h5_fileout) continue if (rank==0 and idx % smi_stride==0): 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) if force_index is not None: spec = spec_data[force_index] rotation = sqr(Umat_data[force_index]) else: spec = spec_data[idx] rotation = sqr(Umat_data[idx]) wavelength_A = np.mean(wave_chans) spectra = iter([(wave_chans, spec, wavelength_A)]) wavlen, flux, wavelength_A = next(spectra) # list of lambdas, list of fluxes, average wavelength assert wavelength_A > 0 assert (len(wavlen)==len(flux)==len(sfall_main)) if np.sum(flux)==0: continue N = crystal.number_of_cells(sfall_main[0].unit_cell()) Ncells_abc = (N,N,N) if not on_axis: Crystal.set_U(rotation) if idx_path is not None: assert( model_file is None) Crystal = utils.open_flex(idx_path)['crystalAB'] if model_file is not None: assert( idx_path is None) P = utils.open_flex(model_file) Crystal = P['crystal'] #mos_spread_deg = P['mos_spread'] #Ncells_abc = P['Ncells_abc'] if force_twocolor: flux *= 0 flux[ilow] = 1e11 flux[ihigh]=1e11 simsAB = sim_utils.sim_twocolors2( Crystal, DET, BEAM, sfall_main, en_chans, flux, pids = None, profile="gauss", oversample=1, Ncells_abc = Ncells_abc, mos_dom=mos_doms, verbose=verbose, mos_spread=mos_spread_deg, cuda=True, device_Id =rank, beamsize_mm=beamsize_mm, exposure_s=exposure_s, accumulate=True, boost=crystal.domains_per_crystal) if not simsAB: continue out = simsAB[0] if add_background: out = out + background if add_noise: SIM = nanoBragg(detector=DET, beam=BEAM) SIM.exposure_s = exposure_s SIM.flux = np.sum(flux) SIM.raw_pixels = flex.double(out.ravel()) 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.add_noise() out = SIM.raw_pixels.as_numpy_array().reshape(out.shape) f = h5py.File(h5_fileout, "w") f.create_dataset("bigsim_d9114", data=out, compression="lzf") ua,ub,uc = Crystal.get_real_space_vectors() f.create_dataset("real_space_a", data=ua) f.create_dataset("real_space_b", data=ub) f.create_dataset("real_space_c", data=uc) f.create_dataset("space_group_hall_symbol", data=Crystal.get_space_group().info().type().hall_symbol()) f.create_dataset("Umatrix", data=Crystal.get_U()) f.create_dataset("fluxes",data=flux) f.create_dataset("energies", data=en_chans) f.close() if npout is not None: np.save(npout, out) # SIM.raw_pixels.as_numpy_array())
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")