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