Commit c575722f authored by Hsing-Yin Chang's avatar Hsing-Yin Chang
Browse files

add execution_fence, use logger for print statements, save output mrc files, etc.

parent 4f89949f
......@@ -3,9 +3,10 @@ import PyNVTX as nvtx
import os
import pygion
from pygion import acquire, attach_hdf5, task, Partition, Region, R, Tunable, WD
from pygion import acquire, attach_hdf5, execution_fence, task, Partition, Region, R, Tunable, WD
from spinifel import settings
from spinifel import settings, utils, contexts, checkpoint
from spinifel.prep import save_mrc
from .prep import get_data
from .autocorrelation import solve_ac
......@@ -17,14 +18,18 @@ from . import mapper
@task(replicable=True)
@nvtx.annotate("legion/main.py", is_prefix=True)
def main():
print("In Legion main", flush=True)
timer = utils.Timer()
total_procs = Tunable.select(Tunable.GLOBAL_PYS).get()
# Reading input images from hdf5
N_images_per_rank = settings.N_images_per_rank
batch_size = min(N_images_per_rank, 100)
max_events = min(settings.N_images_max, total_procs*N_images_per_rank)
writer_rank = 0 # pick writer rank as core 0
# Reading input images using psana2
ds = None
if settings.use_psana:
# For now, we use one smd chunk per node just to keep things simple.
......@@ -32,37 +37,79 @@ def main():
settings.ps_smd_n_events = N_images_per_rank
from psana import DataSource
logger.log("Using psana")
ds = DataSource(exp=settings.exp, run=settings.runnum,
dir=settings.data_dir, batch_size=batch_size,
max_events=max_events)
# Setup logger after knowing the writer rank
logger = utils.Logger(True)
logger.log("In Legion main")
# Load unique set of intensity slices for python process
(pixel_position,
pixel_distance,
pixel_index,
slices, slices_p) = get_data(ds)
logger.log(f"Loaded in {timer.lap():.2f}s.")
solved = solve_ac(0, pixel_position, pixel_distance, slices, slices_p)
logger.log(f"AC recovered in {timer.lap():.2f}s.")
phased = phase(0, solved)
logger.log(f"Problem phased in {timer.lap():.2f}s.")
rho = np.fft.ifftshift(phased.rho_)
print('rho =', rho)
save_mrc(settings.out_dir / f"ac-0.mrc", phased.ac)
save_mrc(settings.out_dir / f"rho-0.mrc", rho)
# Use improvement of cc(prev_rho, cur_rho) to dertemine if we should
# terminate the loop
prev_phased = None
cov_xy = 0
cov_delta = .05
N_generations = settings.N_generations
for generation in range(1, N_generations):
logger.log(f"#"*27)
logger.log(f"##### Generation {generation}/{N_generations} #####")
logger.log(f"#"*27)
# Orientation matching
orientations, orientations_p = match(
phased, slices, slices_p, pixel_position, pixel_distance)
logger.log(f"Orientations matched in {timer.lap():.2f}s.")
# Solve autocorrelation
solved = solve_ac(
generation, pixel_position, pixel_distance, slices, slices_p,
orientations, orientations_p, phased)
logger.log(f"AC recovered in {timer.lap():.2f}s.")
prev_phased = prev_phase(generation, phased, prev_phased)
phased = phase(generation, solved, phased)
logger.log(f"Problem phased in {timer.lap():.2f}s.")
cov_xy, is_cov = cov(prev_phased, phased, cov_xy, cov_delta)
# Check if density converges
if settings.chk_convergence:
cov_xy, is_cov = cov(prev_phased, phased, cov_xy, cov_delta)
if is_cov:
break;
if is_cov:
print("Stopping criteria met!")
break;
rho = np.fft.ifftshift(phased.rho_)
print('rho =', rho)
save_mrc(settings.out_dir / f"ac-{generation}.mrc", phased.ac)
save_mrc(settings.out_dir / f"rho-{generation}.mrc", rho)
execution_fence(block=True)
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