chromatix-team / chromatix

Differentiable wave optics using JAX! Documentation can be found at https://chromatix.readthedocs.io
MIT License
61 stars 12 forks source link

Microscope parameters #20

Closed schneidermc closed 1 month ago

schneidermc commented 1 year ago

Define microscope parameters as part of Microscope class instead of psf.

schneidermc commented 1 year ago

I think we should put some more thought into how to design the simulations, more specifically what parameters and functionality we want to put into the classes Microscope and OpticalSystem. Coming from single-molecule localization microscopy, my view of a microscope is the following: Microscope = Light source (e.g. lasers) + Optical System (e.g. lenses, mirrors, phase masks) + Detector (e.g. camera) Experiment = Sample + Microscope I am happy about any input about other views that would be more suitable for other types of microscopes! Having such a structure and good encapsulation of the parts would make it more intuitive to use chromatix.

diptodip commented 1 year ago

Yes, I think this is still a little bit tricky.

The first thing I'll say is that it would definitely be great to have more realistic sources than the ideal sources we have right now (so modeling things like speckle from a laser or something). And we already have a basic setup for sensors, but it would definitely be nice if they would be compatible with the Microscope class as well. So that can be something we rework a little bit which shouldn't be too hard.

The second thing is, do take a look at the new Microscope class if you haven't already! I think it solves at least the initial proposal of this issue (defining Microscope parameters in the Microscope class).

As for describing systems, there are two ways we could build one right now. One way involves including everything you talked about, including source, optical elements, sample, and sensor, in an OpticalSystem. Then, we are simulating Fields passing from the light source, through the system and sample, and then finally reaching the sensor. This is very flexible in that you can include whatever you want in any order and you will still get a Field out. To me, this is also a pretty natural way to describe a system (including the sample). For example, this style lets you do something like optimizing the sample in a convenient way, so this is how someone doing phase retrieval might want to describe their system.

But since this description is inefficient if we know that the system is linear and shift invariant, we have the Microscope class for systems with a shift invariant PSF. This lets us just describe the PSF and then simulate the imaging with a convolution, which is the only real reason that we have a Microscope class at all. If it weren't more efficient, we could just have OpticalSystem. This kind of description is more useful if you want to optimize some parameter of the optical elements, like the pixels of the phase mask in my case.

The current version of the Microscope wants a Module/function to describe the PSF of the system. If you look at the very last cell of the Holoscope example you'll see what it looks like right now to define a Microscope with a 4f system. We could also imagine other kinds of systems that are rotationally invariant instead of spatially invariant, or systems with field varying PSFs, etc. And we could come up with versions of the Microscope that would handle the simulation of imaging from such systems in the most efficient way.

So right now, the Microscope can still be useful for any other kind of PSF you would want to describe or for PSFs that you would like to parameterize in different ways, but it still isn't this natural object-oriented hierarchy that you are describing.

The tricky part is that if we tried to do something like what you are suggesting by just providing a list of what elements are in the microscope, we would need some way to choose for each element what actual arguments need to be passed to construct each element (e.g. a lens needs the microscope objective's focal length, refractive index, and numerical aperture but a phase mask would also need the phase). There are also other details like calculating the spacing required for the PSF simulation. Right now, these details can be taken care of by a PSF module for different kinds of systems, where the PSF module also knows about the Microscope and its details (f, n, NA, spacing, etc.).

Sorry for the long explanation but I'd love to know what you think!

diptodip commented 1 year ago

Working on this in PR #92, where we have added the ability to use sensors in Microscopes as well as in OpticalSystems. So while before you could have already placed a sensor at the end of an OpticalSystem to measure intensity of the incoming Field, now we can also use them in a Microscope like so:

holoscope = Microscope(
    system_psf=Optical4FSystemPSF(
        shape=(1280, 1280),
        spacing=0.325,
        phase=trainable(partial(defocused_ramps, delta=[1582.0] * 6))
    ),
    sensor=BasicShotNoiseSensor(
        shape=(256, 256),
        spacing=1.625,
        shot_noise_mode="poisson",
        resampling_method="cubic",
        reduce_axis=0 # if None, returns a batch of 2D images
    ),
    f=100.0,
    n=1.33,
    NA=0.8,
    spectrum=0.532,
    spectral_density=1.0,
    padding_ratio=0.5,
    taper_width=5.0,
)

This Microscope still defines the PSF of the system instead of just defining the system, but like we talked about that's the simplest way we can keep this computation efficient and flexible for other systems.

schneidermc commented 1 year ago

It's great that we have sensors now! Regarding the naming, I think it is a bit counterintuitive that the BasicShotNoiseSensor defaults to shot_noise_mode=None, i.e. no shot noise. Maybe we want to change the name to BasicSensor, BasicCamera or just Camera, depending on what other sensors or detectors we potentially might want to add in the future.