Skip to content

Conversation

@michaelxu01
Copy link
Contributor

working for conventional and gradient, but maybe not the most elegant for the gradient engine. open to suggestions

@michaelxu01
Copy link
Contributor Author

I'm also not a huge fan of default value of 0

Copy link
Owner

@hexane360 hexane360 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks pretty good, made a couple of notes. Why don't you rename the key to 'pos_update_mag' or 'pos_update_rms' to make clear it pertains to the position solver? Do we want 'tilt_update_mag'/'tilt_update_rms' as well?

state, iter_solver_states[sol_i], filter_vars(iter_grads, solver.params), losses['total_loss']
)
state = apply_update(state, update)
update_mag = xp.linalg.norm(update['positions'], axis=-1)
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What happens if no position solver is specified, and 'positions' is not in update?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should be addressed now

Comment on lines 294 to 296
for (k) in other_keys:
progress[k].iters.append(i + start_i)
progress[k].values.append(xp.mean(update_mag))
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The for loop seems redundant if you're special-casing 'update_mag' regardless

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah I realize it'll definitely break, hopefully fixed in newest commit

@michaelxu01 michaelxu01 requested a review from hexane360 October 31, 2025 19:06
other_keys = ('positions_update_rms', ) if position_solver is not None else ()
# populate missing keys in progress dictionary
for k in ('detector_loss', 'total_loss'):
for k in (*('detector_loss', 'total_loss'), *other_keys):
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
for k in (*('detector_loss', 'total_loss'), *other_keys):
for k in ('detector_loss', 'total_loss', *other_keys):


# check positions are at least overlapping object
sim.state.object.sampling.check_scan(sim.state.scan, sim.state.probe.sampling.extent / 2.)
# for k in ('update_mag', ):
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
# for k in ('update_mag', ):

Comment on lines 113 to 114
pos_update_rms = xp.linalg.norm(pos_update, axis=-1, keepdims=True)
logger.info(f"Position update: mean {xp.mean(pos_update_rms)}")
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
pos_update_rms = xp.linalg.norm(pos_update, axis=-1, keepdims=True)
logger.info(f"Position update: mean {xp.mean(pos_update_rms)}")
pos_update_rms = float(xp.mean(xp.linalg.norm(pos_update, axis=-1, keepdims=True)))
logger.info(f"Position update: mean {pos_update_rms}")

Comment on lines 121 to 122
progress['positions_update_rms'].iters.append(i + start_i)
progress['positions_update_rms'].values.append(xp.mean(pos_update_rms))
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
progress['positions_update_rms'].iters.append(i + start_i)
progress['positions_update_rms'].values.append(xp.mean(pos_update_rms))
progress['positions_update_rms'].iters.append(i + start_i)
progress['positions_update_rms'].values.append(pos_update_rms)

michaelxu01 and others added 2 commits October 31, 2025 21:12
Co-authored-by: Colin Gilgenbach <colin@gilgenbach.net>
Co-authored-by: Colin Gilgenbach <colin@gilgenbach.net>
@hexane360 hexane360 merged commit 1fd227c into hexane360:develop Nov 1, 2025
4 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants