Commit 396cd1d4 authored by chuckie82's avatar chuckie82
Browse files

set pickle protocol to 4 to match lcls

parent 6902b35b
......@@ -11,7 +11,7 @@ def generate_checkpoint_name(outdir: str, gen_num: int, tag='') -> str:
return os.path.join(outdir, fname)
def save_checkpoint(res: dict, outdir: str, gen_num: int, tag='') -> None:
def save_checkpoint(res: dict, outdir: str, gen_num: int, tag='', protocol=4) -> None:
"""Save results as pickle"""
# create outdir if does not exist
Path(outdir).mkdir(parents=True, exist_ok=True)
......@@ -19,7 +19,7 @@ def save_checkpoint(res: dict, outdir: str, gen_num: int, tag='') -> None:
# save pickle
fname = generate_checkpoint_name(outdir, gen_num, tag)
with open(fname, 'wb') as handle:
pickle.dump(res, handle, protocol=pickle.HIGHEST_PROTOCOL)
pickle.dump(res, handle, protocol=protocol)
def load_checkpoint(outdir: str, gen_num: int, tag='') -> dict:
......
......@@ -90,7 +90,7 @@ def main():
'slices_': slices_,
'ac': ac
}
checkpoint.save_checkpoint(myRes, settings.out_dir, curr_gen, tag="solve_ac")
checkpoint.save_checkpoint(myRes, settings.out_dir, curr_gen, tag="solve_ac", protocol=4)
ac_phased, support_, rho_ = phase(curr_gen, ac)
logger.log(f"Problem phased in {timer.lap():.2f}s.")
......@@ -101,7 +101,7 @@ def main():
'support_': support_,
'rho_': rho_
}
checkpoint.save_checkpoint(myRes, settings.out_dir, curr_gen, tag="phase")
checkpoint.save_checkpoint(myRes, settings.out_dir, curr_gen, tag="phase",protocol=4)
# Save electron density and intensity
rho = np.fft.ifftshift(rho_)
intensity = np.fft.ifftshift(np.abs(np.fft.fftshift(ac_phased)**2))
......@@ -131,7 +131,7 @@ def main():
'pixel_distance_reciprocal': pixel_distance_reciprocal,
'orientations': orientations
}
checkpoint.save_checkpoint(myRes, settings.out_dir, generation, tag="match")
checkpoint.save_checkpoint(myRes, settings.out_dir, generation, tag="match",protocol=4)
# Solve autocorrelation
ac = solve_ac(
......@@ -147,7 +147,7 @@ def main():
'ac_phased': ac_phased,
'ac': ac
}
checkpoint.save_checkpoint(myRes, settings.out_dir, generation, tag="solve_ac")
checkpoint.save_checkpoint(myRes, settings.out_dir, generation, tag="solve_ac",protocol=4)
if comm.rank == 0:
prev_rho_ = rho_[:]
......@@ -163,7 +163,7 @@ def main():
'support_': support_,
'rho_': rho_
}
checkpoint.save_checkpoint(myRes, settings.out_dir, generation, tag="phase")
checkpoint.save_checkpoint(myRes, settings.out_dir, generation, tag="phase",protocol=4)
# Check if density converges
......@@ -198,7 +198,7 @@ def main():
'rho_': rho_,
'orientations': orientations
}
checkpoint.save_checkpoint(myRes, settings.out_dir, generation)
checkpoint.save_checkpoint(myRes, settings.out_dir, generation, tag="", protocol=4)
logger.log(f"Results saved in {settings.out_dir}")
logger.log(f"Successfully completed in {timer.total():.2f}s.")
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment