Skip to content

Lower memory in photon to geometry calculation#515

Draft
welucas2 wants to merge 3 commits intomainfrom
u/welucas2/diffraction-memory-usage
Draft

Lower memory in photon to geometry calculation#515
welucas2 wants to merge 3 commits intomainfrom
u/welucas2/diffraction-memory-usage

Conversation

@welucas2
Copy link
Copy Markdown
Collaborator

With Jim's investigation this may no longer be necessary, but in case it's useful here are the changes I made to diffraction.py that reduce memory usage in directed_dist(). The code as it is now in main is nicely vectorised, but relies heavily on intermediate arrays which can really eat into memory if bright objects are being processed. This line alone allocates very nearly 1 GB when I run the realistic test Jim gave me with photon pooling enabled and photon processing of bright objects.

(That's why it may not be needed - if it's only ever going to be FFTs.)

I've split the calculation apart so it no longer calculates and stores the vectors between all points (photons) and structures, but it now stores only the vector to the nearest object.

I also investigated changing the call to apply_diffraction_delta() in e.g. RubinDiffraction by having

imSim/imsim/photon_ops.py

Lines 294 to 301 in 6dda08b

v = self.apply_diffraction_delta(
np.c_[x, y],
v,
photon_array.time,
wavelength=photon_array.wavelength * 1.0e-9,
geometry=RUBIN_SPIDER_GEOMETRY,
distribution=self.diffraction_rng(rng),
)

not call it all in one go, but looping through the full array of v and filling it in chunks of say 100. I'm not sure it actually helped much at all in contrast with this change and the one in #513, so I haven't included it here, but wanted to mention it in case.

…ctures in geometry rather than vectorising across large arrays.
…. Not as clean or Pythonic, but it saves a *lot* of memory.
@welucas2 welucas2 changed the title U/welucas2/diffraction memory usage Lower memory in photon to geometry calculation Mar 27, 2026
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.

1 participant