def sim_all(Nshot_max=None, jid=0, n_jobs=1):
odir = "/global/project/projectdirs/lcls/dermen/bigsim"
prefix = odir + "/run62_%06d"
spec_file = h5py.File(full_path("simMe_data_run62.h5"), "r")
spec_data = spec_file["hist_spec"]
Umat_data = spec_file["Umats"]
en_chans = spec_file["energy_bins"][()]
wave_chans = ENERGY_CONV / en_chans
sfall = sim_spectra.load_spectra(full_path("test_sfall.h5"))
# his jobs id
#jid = int(sys.argv[1])
#global device_Id
device_Id = jid
Nshot = spec_data.shape[0]
if Nshot_max is not None:
Nshot = min(Nshot_max, Nshot)
print("Job %d: Simulating %d shots" % (jid, Nshot))
# total number of available jobs
#n_jobs = int(sys.argv[2])
idx_range = np.array_split(np.arange(Nshot), n_jobs)
Nshot_per_job = len(idx_range[jid])
for idx in idx_range[jid]:
print("<><><><><><>")
print("Job %d; Image %d / %d" % (jid, idx + 1, Nshot_per_job))
print("<><><><><><>")
if (jid == 0):
print("GPU status")
os.system("nvidia-smi")
spec = spec_data[idx]
C_ori = sqr(Umat_data[idx])
wavelength_A = np.mean(wave_chans)
C = microcrystal(Deff_A=Deff_A,
length_um=length_um,
beam_diameter_um=beam_size_mm * 1000,
verbose=False)
iterator = iter([(wave_chans, spec, wavelength_A)])
run_sim2smv(
prefix=prefix % idx, # "run62_%06d" % idx,
crystal=C,
spectra=iterator,
sfall_main=sfall,
rotation=C_ori,
quick=False,
rank=jid,
save_bragg=False)
def tst_one(quick=False, prefix="step5"):
spec_file = h5py.File(full_path("test_data.h5"), "r")
idx = 2
en_chans = spec_file["energy_bins"][()]
wave_chans = ENERGY_CONV / en_chans
spec = spec_file["hist_spec"][idx]
C_ori = sqr(spec_file["Umats"][idx])
wavelength_A = np.mean(wave_chans)
sfall = sim_spectra.load_spectra(full_path("test_sfall.h5"))
C = microcrystal(Deff_A=Deff_A,
length_um=length_um,
beam_diameter_um=beam_size_mm * 1000)
iterator = iter([(wave_chans, spec, wavelength_A)])
run_sim2smv(prefix="tst_one_test_data_%06d" % idx,
crystal=C,
spectra=iterator,
sfall_main=sfall,
rotation=C_ori,
quick=False,
rank=0,
save_bragg=False) #, actually_use_smv=False)
parser.add_argument("-o",
dest='o',
help='outputfile prefix (prefix.npz) ',
type=str,
default=None)
parser.add_argument("-plot",
dest='plot',
action='store_true',
help='plot Ymodel vs Yobs')
parser.add_argument("-anom",
dest='anom',
action='store_true',
help='group h by anom')
args = parser.parse_args()
SFall = sim_spectra.load_spectra("test_sfall.h5")
# need to load the structure factors at the channel A and B wavelengths
spec_f = h5py.File("simMe_data_run62.h5", "r")
en_bins = spec_f["energy_bins"][()]
ilow = abs(en_bins - ENERGY_LOW).argmin()
ihigh = abs(en_bins - ENERGY_HIGH).argmin()
SA = SFall[ilow].amplitudes() # energy A (lower energy)
SB = SFall[ihigh].amplitudes() # energy B (higher energy)
# load the integrated data
df = pandas.read_pickle(args.i)
Kfact = 1e20
if args.anom:
hkl = tuple(map(tuple, df[['hAnom', 'kAnom', 'lAnom']].values.astype(int)))
def run_sim2smv(Nshot_max, odir, prefix, rank, n_jobs, save_bragg=False,
save_smv=True, save_h5 =False, return_pixels=False):
from six.moves import range, StringIO
from six.moves import cPickle as pickle
from cxid9114.sim import sim_utils
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.parameters import ENERGY_CONV, ENERGY_HIGH, ENERGY_LOW
from cxid9114.bigsim import sim_spectra
odir_j = os.path.join( odir, "job%d" % rank)
if not os.path.exists(odir_j):
os.makedirs(odir_j)
add_noise = False
add_background = False
overwrite = True #$False
sample_thick_mm = 0.005 # 50 micron GDVN nozzle makes a ~5ish micron jet
air_thick_mm =0 # mostly vacuum, maybe helium layer of 1 micron
flux_ave=2e11
add_spots_algorithm="cuda"
big_data = "." # directory location for reference files
detpixels_slowfast = (1800,1800)
pixsize_mm=0.11
distance_mm = 125
offset_adu=30
mos_spread_deg=0.015
mos_doms=1000
beam_size_mm=0.001
exposure_s=1
use_microcrystal=True #False
Ncells_abc=(120,120,120)
Deff_A = 2200
length_um = 2.2
timelog = False
background = utils.open_flex("background")
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, prefix % idx + ".img")
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)
spec = spec_data[2]
rotation = sqr(Umat_data[2])
wavelength_A = np.mean(wave_chans)
spectra = iter([(wave_chans, spec, wavelength_A)])
direct_algo_res_limit = 1.7
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))
N = crystal.number_of_cells(sfall_main[0].unit_cell())
if use_mcrocrystal:
Ncells_abc = (N,N,N)
flux *= 0
flux[ilow] = 1e12
flux[ihigh]=1e12
UMAT_nm = flex.mat3_double()
mersenne_twister = flex.mersenne_twister(seed=0)
scitbx.random.set_random_seed(1234)
rand_norm = scitbx.random.normal_distribution(mean=0, sigma=mos_spread_deg * math.pi/180.)
g = scitbx.random.variate(rand_norm)
mosaic_rotation = g(mos_doms)
for m in mosaic_rotation:
site = col(mersenne_twister.random_double_point_on_sphere())
UMAT_nm.append( site.axis_and_angle_as_r3_rotation_matrix(m,deg=False) )
Amatrix_rot = (rotation * sqr(sfall_main[0].unit_cell().orthogonalization_matrix())).transpose()
#SIM = nanoBragg(detpixels_slowfast=detpixels_slowfast,
# pixel_size_mm=pixsize_mm,Ncells_abc=Ncells_abc,
# wavelength_A=wavelength_A,verbose=verbose)
SIM = nanoBragg(detector=DET, beam=BEAM, panel_id=0,verbose=verbose)
SIM.adc_offset_adu = offset_adu # Do not offset by 40
SIM.mosaic_spread_deg = mos_spread_deg # interpreted by UMAT_nm as a half-width stddev
SIM.mosaic_domains = mos_doms
SIM.distance_mm=distance_mm
SIM.set_mosaic_blocks(UMAT_nm)
SIM.seed = 1
SIM.oversample=1
SIM.wavelength_A = wavelength_A
SIM.polarization=1
SIM.default_F=0
SIM.Fhkl=sfall_main[0].amplitudes()
SIM.progress_meter=False
SIM.flux=flux_ave
SIM.exposure_s=exposure_s
SIM.beamsize_mm=beam_size_mm
SIM.Ncells_abc=Ncells_abc
SIM.xtal_shape=shapetype.Gauss
idxpath = "try3_idx2/job0/dump_0_data.pkl"
Amat = sim_utils.Amatrix_dials2nanoBragg(utils.open_flex(idxpath)["crystalAB"])
#SIM.Umatrix = rotation.elems
#SIM.unit_cell_Adeg = sfall_main[0].unit_cell()
#SIM2 = nanoBragg(detector=DET, beam=BEAM, panel_id=0,verbose=verbose)
#SIM2.adc_offset_adu = offset_adu # Do not offset by 40
#SIM2.mosaic_spread_deg = mos_spread_deg # interpreted by UMAT_nm as a half-width stddev
#SIM2.mosaic_domains = mos_doms
#SIM2.distance_mm=distance_mm
#SIM2.set_mosaic_blocks(UMAT_nm)
#SIM2.seed = 1
#SIM2.oversample=1
#SIM2.wavelength_A = wavelength_A
#SIM2.polarization=1
#SIM2.default_F=0
#SIM2.Fhkl=sfall_main[0].amplitudes()
#SIM2.progress_meter=False
#SIM2.flux=flux_ave
#SIM2.exposure_s=exposure_s
#SIM2.beamsize_mm=beam_size_mm
#SIM2.Ncells_abc=Ncells_abc
#SIM2.xtal_shape=shapetype.Gauss
#SIM2.Umatrix = rotation.elems
#SIM2.unit_cell_Adeg = sfall_main[0].unit_cell()
#SIM.Amatrix_RUB = Amatrix_rot
#SIM.Amatrix = Amatrix_rot.inverse()
raw_pixel_sum = flex.double(len(SIM.raw_pixels))
Nflux = len(flux)
print("Beginning the loop")
for x in range(Nflux):
if x % 10==0:
print("+++++++++++++++++++++++++++++++++++++++ Wavelength %d / %d" % (x+1, Nflux) , end="\r")
if flux[x] ==0:
continue
#SIM.Amatrix = Amatrix_rot.inverse()
print (SIM.Ncells_abc)
SIM.wavelength_A=wavlen[x]
SIM.flux=flux[x]
SIM.Fhkl=sfall_main[x].amplitudes()
SIM.Amatrix = Amat#Amatrix_rot.inverse()
#sim_utils.compare_sims(SIM, SIM2)
#SIM.Ncells_abc=Ncells_abc
#SIM.adc_offset_adu = offset_adu
#SIM.mosaic_spread_deg = mos_spread_deg # interpreted by UMAT_nm as a half-width stddev
#SIM.mosaic_domains = mos_doms #
#SIM.distance_mm=distance_mm
#SIM.set_mosaic_blocks(UMAT_nm)
#SIM.seed = 1
#SIM.polarization=1
#SIM.default_F=0
#SIM.xtal_shape=shapetype.Gauss
#SIM.progress_meter=False
#SIM.exposure_s = exposure_s
#SIM.beamsize_mm=beam_size_mm
SIM.timelog=timelog
SIM.device_Id=rank
SIM.raw_pixels *= 0 # just in case!
SIM.add_nanoBragg_spots_cuda()
if use_microcrystal:
raw_pixel_sum += SIM.raw_pixels * crystal.domains_per_crystal
print()
SIM.raw_pixels = raw_pixel_sum
if add_background:
SIM.raw_pixels = SIM.raw_pixels + background
SIM.detector_psf_kernel_radius_pixels=5;
#SIM.detector_psf_fwhm_mm=0.08;
#SIM.detector_psf_type=shapetype.Fiber # rayonix=Fiber, CSPAD=None (or small Gaussian)
SIM.detector_psf_type=shapetype.Unknown # for CSPAD
SIM.detector_psf_fwhm_mm=0
SIM.quantum_gain = 28.
#SIM.apply_psf()
if add_noise:
SIM.add_noise() #converts phtons to ADU.
extra = "PREFIX=%s;\nRANK=%d;\n"%(prefix,rank)
out = SIM.raw_pixels.as_numpy_array()
if save_smv:
SIM.to_smv_format_py(fileout=smv_fileout,intfile_scale=1,rotmat=True,extra=extra,gz=True)
elif save_h5:
f = h5py.File(h5_fileout, "w")
f.create_dataset("bigsim_d9114",
data=SIM.raw_pixels.as_numpy_array().astype(np.uint16).reshape(detpixels_slowfast),
compression="lzf")
f.close()
if npout is not None:
np.save(npout, SIM.raw_pixels.as_numpy_array())
SIM.free_all()
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")
parser.add_argument('--filt', action='store_true', dest='filt')
parser.add_argument('--gauss', action='store_true')
args = parser.parse_args()
import pandas
from joblib import Parallel, delayed
import numpy as np
from itertools import izip
import sys
#from cxid9114 import utils
from cxid9114.bigsim import sim_spectra
#from cxid9114.parameters import ENERGY_HIGH, ENERGY_LOW
import h5py
from cxid9114.solvers.filter_outliers import filter_outliers
spec = sim_spectra.load_spectra("bigsim_rocketships/test_sfall.h5")
spec_f = h5py.File("bigsim_rocketships/test_data.h5", "r")
en_bins = spec_f["energy_bins"][()]
ilow = 19 # abs(en_bins - ENERGY_LOW).argmin()
ihigh = 110 #abs(en_bins - ENERGY_HIGH).argmin()
SA = spec[ilow].amplitudes() # energy A (lower energy)
SB = spec[ihigh].amplitudes() # energy B (higher energy)
# load the integrated data
df = pandas.read_pickle(args.i)
df.reset_index(inplace=True)
if args.hanom:
hkey = ['hAnom', 'kAnom', 'lAnom']
else:
def run_paramList(Ntrials, odir, tag, rank, n_jobs, pkl_file):
import os
import sys
from copy import deepcopy
import numpy as np
import h5py
from IPython import embed
from scipy.ndimage.morphology import binary_dilation
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 dxtbx.model.crystal import CrystalFactory
from cctbx import crystal,crystal_orientation
from cxid9114 import utils
from cxid9114.sim import sim_utils
from cxid9114.spots import spot_utils
from cxid9114.bigsim.bigsim_geom import DET,BEAM
from cxid9114.parameters import ENERGY_CONV, ENERGY_HIGH, ENERGY_LOW
from cxid9114.refine.jitter_refine import make_param_list
from cxid9114.bigsim import sim_spectra
from LS49.sim.step4_pad import microcrystal
data_pack = utils.open_flex(pkl_file)
CRYST = data_pack['crystalAB']
mos_spread_deg=0.015
mos_doms=1000
beam_size_mm=0.001
exposure_s=1
use_microcrystal=True
Deff_A = 2200
length_um = 2.2
timelog = False
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")
refls_strong = data_pack['refls_strong']
strong_mask_img = spot_utils.strong_spot_mask(
refls_strong, (1800,1800) )
# edge detection in the ground truth strong mask image
reference_img = (binary_dilation(strong_mask_img, iterations=1).astype(int) -
strong_mask_img.astype(int) ).astype(bool)
param_fileout = os.path.join( odir, "rank%d_%s.pkl" % (rank, tag))
param_list = make_param_list(
CRYST, DET, BEAM, Ntrials,
rot=0.09, cell=0.1, eq=(1,1,0),
min_Ncell=20, max_Ncell=40,
min_mos_spread=0.005, max_mos_spread=0.02)
for p in param_list:
print(p['crystal'].get_unit_cell().parameters())
shot_idx = int(data_pack['img_f'].split("_")[-1].split(".")[0])
Fluxes = spec_data[2] #shot_idx]
Pmax = param_list[0]
F1max = 0
for i_trial in range(Ntrials):
print ("<><><><><><><><><><><><><><>")
print ("Job %d; Trial %d / %d" % (rank, i_trial+1, Ntrials))
print ("<><><><><><><><><><><><><><>")
if (rank==0 and i_trial % 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)
assert (len(wave_chans)==len(Fluxes)==len(sfall_main))
if np.sum(Fluxes)==0:
print ("Cannot simulate with an all-zeros spectrum!")
sys.exit()
N = crystal.number_of_cells(sfall_main[0].unit_cell())
Ncells_abc = (N,N,N)
if force_twocolor:
Fluxes *= 0
Fluxes[ilow] = 1e12
Fluxes[ihigh]=1e12
P = param_list[i_trial]
simsAB = sim_utils.sim_twocolors2(
P['crystal'],
DET,
BEAM,
sfall_main,
en_chans,
Fluxes,
pids = None,
profile="gauss",
oversample=0,
Ncells_abc = Ncells_abc, #P['Ncells_abc'],
mos_dom=mos_doms,
verbose=verbose,
mos_spread=mos_spread_deg,
#@mos_spread=P['mos_spread'],
cuda=True,
device_Id=rank,
beamsize_mm=beamsize_mm,
exposure_s=exposure_s,
boost=crystal.domains_per_crystal)
out = np.sum( [ simsAB[i][0] for i in simsAB.keys() if simsAB[i]], axis=0)
if out.shape==():
print("This simsAB output has an empty shape, something is wrong!")
sys.exit()
trial_refls = spot_utils.refls_from_sims([out], DET, BEAM, thresh=thresh)
trial_spotmask = spot_utils.strong_spot_mask(
trial_refls, (1800,1800) )
trial_img = (binary_dilation(trial_spotmask, iterations=1).astype(int) -
trial_spotmask.astype(int) ).astype(bool)
comp_img = trial_img.astype(int) + reference_img.astype(int)*2
Nfalse_neg = (comp_img==2).sum() # reference has signal but trial doesnt
Nfalse_pos = (comp_img==1).sum() # trial has signal but reference doesnt
Ntrue_pos = (comp_img==3).sum() #
Precision = float(Ntrue_pos) / (Ntrue_pos + Nfalse_pos)
Recall = float(Ntrue_pos) / (Ntrue_pos + Nfalse_neg)
F1 = 2.*(Precision*Recall) / (Precision+Recall)
if F1 > F1max:
Pmax = {'crystal': P['crystal'],
'mos_spread': P['mos_spread'],
'Ncells_abc': P['Ncells_abc'], "F1": F1}
F1max = F1
print("Rank %d, Trial %d: F1score = %.5f" % (rank, i_trial, F1))
utils.save_flex(Pmax, param_fileout)
return F1max
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")