GalSim-developers / GalSim

The modular galaxy image simulation toolkit. Documentation:
http://galsim-developers.github.io/GalSim/
Other
224 stars 106 forks source link

Construct PhotonArray from pre-allocated memory #1108

Closed jmeyers314 closed 3 years ago

jmeyers314 commented 3 years ago

I've been thinking a bit more about how we might batch photons together to efficiently raytrace them using batoid. I think it'd be useful to be able to construct a PhotonArray from (portions of) previously allocated numpy arrays. E.g.,


# allocate in shared memory...
N = 10_000
x = np.zeros(N)
y = np.zeros(N)
flux = np.zeros(N)
dxdz = np.zeros(N)
dydz = np.zeros(N)
wavelength = np.zeros(N)

pa_batch = galsim.PhotonArray._fromArrays(x, y, flux, dxdz, dydz, wavelength)
pa1 = galsim.PhotonArray._fromArrays(
    x[:N//2],
    y[:N//2],
    flux[:N//2],
    dxdz[:N//2],
    dydz[:N//2],
    wavelength[:N//2]
)

pa2 = galsim.PhotonArray._fromArrays(
    x[N//2:],
    y[N//2:],
    flux[N//2:],
    dxdz[N//2:],
    dydz[N//2:],
    wavelength[N//2:]
)

# in process 1
obj1._shoot(pa1, rng1)
# simultaneously in process 2
obj2._shoot(pa2, rng2)
# apply photon operations to entire batch without doing any copying...
batoidPhotonOp.applyTo(pa_batch)
# accumulate into sensor...

I.e., with PhotonArray._fromArrays one could perform the steps in drawImage, using multiprocessing to fill in independent objects photons "at the entrance pupil", then use batoid to trace them all simultaneously, and then finish with accumulate. In fact, I think most of our photonOps, and even the atmospheric kicks, could be done simultaneously, not just batoid. (Though DCR probably doesn't work, since that requires a local_wcs which will be different for different objects). I'm not sure how to integrate this into drawImage() itself (probably requires some kind of delayed-execution, which would be tricky), but this at least provides a hook to drawing things "manually".

rmjarvis commented 3 years ago

This sounds very reasonable. My one quick API push back would be to not make the function start with an underscore. There's no reason for this to be a "private" method. It should be a normal classmethod in good standing for users to use and expect good documentation and some version stability, IMO. So I think it should just be named fromArrays.

jmeyers314 commented 3 years ago

My thought for making this "private" was that I didn't want to necessarily spend cycles checking if the arrays are all the same length, dtype=float, contiguous, etc. I think a public interface probably should check those things. One solution would be to implement both, of course.

rmjarvis commented 3 years ago

Sure. Could do our normal thing of having the public version do the checks and then call the private version that doesn't. And I kind of consider these "leading underscore" versions with otherwise the same name to be part of the official API. As opposed to functions that start with an underscore without a corresponding public function, which we use as implementation details.