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")