diff --git a/phaser/engines/conventional/run.py b/phaser/engines/conventional/run.py index 2264805..14ce434 100644 --- a/phaser/engines/conventional/run.py +++ b/phaser/engines/conventional/run.py @@ -63,6 +63,14 @@ def run_engine(args: EngineArgs, props: ConventionalEnginePlan) -> ReconsState: # save progress, it will get clobbered by JIT kernels progress = sim.state.progress + # populate missing keys in progress dictionary + for k in ('detector_loss', 'total_loss'): + if k not in sim.state.progress: + sim.state.progress[k] = ProgressState() + + # save progress, it will get clobbered by JIT kernels + progress = sim.state.progress + observer.init_engine( sim.state, recons_name=args['recons_name'], plan=props, noise_model=noise_model.name(), diff --git a/phaser/engines/gradient/run.py b/phaser/engines/gradient/run.py index 7482f1d..2b44369 100644 --- a/phaser/engines/gradient/run.py +++ b/phaser/engines/gradient/run.py @@ -277,8 +277,8 @@ def run_engine(args: EngineArgs, props: GradientEnginePlan) -> ReconsState: # report losses normalized by # of probe positions losses = tree.map(lambda v: float(v / groups.n_pos), losses) for (k, v) in losses.items(): - progress[k].iters.append(i + start_i) - progress[k].values.append(v) + progress[k].iters.append(int(i + start_i)) + progress[k].values.append(float(v)) # update per-iteration solvers for (sol_i, solver) in enumerate(iter_solvers): diff --git a/phaser/execute.py b/phaser/execute.py index dc7c8a2..dfcbb06 100644 --- a/phaser/execute.py +++ b/phaser/execute.py @@ -233,6 +233,9 @@ def initialize_reconstruction( ) -> PreparedRecons: logging.basicConfig(level=logging.INFO) + with open("post_init_plan.json", "w") as f: + pane.write_json(plan, f, indent=4) + if xp is not None: xp = cast_array_module(xp) # TODO: nicer output here