Source and Scene: High-level science classes

[LouisDesdoigts]

So after our big science planning zoom call on Friday I now have a much better idea of where we want to take the package for JWST. I believe we would be best served by creating two new class types:

---

Source: This class would contain and parameterise the astrophysical source object, and determine the control flow of its modelling. Ie there are two (semi) distinct methods used to model an unresolved point source versus a resolved source with an extended distribution (proto-planetary disks, etc). Similarly there are multiple levels of detail we could want here. Extended objects over a large field of view may need to be modeled anisoplantaically (changing with on-sky position) depending on the optical system. This can actually complicate our life a decent amount with a term I'm going to call 'the curse of the forwards model'. If we have multiple source objects, some resolved and some unresolved, within a small on-sky separation then we can no longer apply the same computational control flow for an entire detector/image. This is where the Source object would allow each specific set of objects call different sub-routines based on their type. Ie we may want a different regularisation term for proto-planetary disks vs a Wolf-Rayet star.

At this point I can only loosely suggest some class attributes without knowing how we will integrate this into the wider ∂Lux infrastructure, but I suspect @benjaminpope would be the most knowledgeable on what we would want here. Off the top of my head:

• Spectrum: array, (nwavels, 2), Wavelengths and weights
• resolved: bool
• type: str, ('point', 'extended')
• regularisation: 'str' ('max_entropy', ... )
• position: array (RA, DEC)

This class would slightly change the OpticalSystem.weights class attribute, transforming it into its more physical interpretation, the filter. Similarly it would allow us to choose the computational cost of our models by having low-importance objects use degraded source spectrums that use a fewer number of wavelengths than the science target. Similarly we could specify oversampling rates for important targets etc.

---

Scene: This would store a list of of the Source objects and determine the control flow for how each is modeled (ie image plane convolution for extended sources, current methods for unresolved sources). I envision this class containing the code used to perform convolutions, regularisation, parametric transit models (maybe?) - you get the idea, all the not-psf modelling code needed to create a final image or parameterise any number of external sources. This class would call the OpticalSystem to generate psfs and then would handle pasting together the psfs correctly (say we model some with oversampling and some without), and passing them to a potential detector class that maybe will be teased out of the the current OpticalSystem (more info on this below).

One thing I will say though is that I really hate the term 'Scene', I'm not entirely sure why but @Jordan-Dennis agrees with me 😂. I think it feels way too amorphous, but like AstrophysicalScene is too wordy etc... Open to any and all suggestions!

---

I think this is a really important step as this will begin to be where we really start to become not just a better PSF modeling engine than poppy based on the same principles, but a full end-to-end astrophysical autodiff forwards modelling framework. We should, in theory, be able to model ANY set of data from ANY telescope with any arbitrary set of sources and parameterisations of telescopes, transit models, regularisations etc.

These additions would open some questions too about what code lives where - Does the distortion mapping happen in the OpticalSystem or the Scene object. If the Source object contains the RA, DEC information and the OpticalSystem contains the positions attribute how do we coherently mix these two together - perhaps a base optical system class that only contains the core functionality to be used by the Scene object. The current OpticalSystem would then inherit from this, and we would use a new class type that it inherits from that doesn't contain positions, bur rather pointing in RA, DEC - helping drive the distortion mapping etc.

This does also partly open the question as to whether maybe we should separate the detector modelling code (which is pretty basic at this point) into its own class type to store distortion functions and what not that can be driven by any of the different class types.

---

Anyway this is the grand vision that I have in my head, I'm very open to suggestions and am aware that we probably only want to build a very reduced functionality version to do our specific science first. Having a coherent wider picture of how all this will paste together I think will save us a lot of headaches in the long run. I envision the dLux/base.py script then containing these classes:

• OpticalSystem (or the Base optical system and the classes that inherit from it)
• Source
• Scene
• Detector (maybe, as discussed above)

Let me know if anything isn't clear here or whatever thoughts you may have!

[benjaminpope]

Hi Louis,

Just to get some points in there... probably more thoughts in due course...

Anisoplanatism

Direction-dependent aberrations are in general hard to deal with, but I think it is not as bad as you fear! Resolved structure like disks etc will never be over sufficiently wide fields of view that this matters. The case where we will deal with wide fields of view is with point sources - we may sometimes look at wide-ish binaries (eg some of the HST brown dwarfs), or at fields of many stars.

But in general, I think high angular resolution imaging will only encounter a single set of aberrations, where the use case is simultaneous phase retrieval and image deconvolution / point source astrometry.

On the other hand, where we will want to use pointing dependent aberrations is in wide angle astrometry (Toliman, or other foreseeable stuff with JWST) or pure phase retrieval (HST/JWST technical stuff)

Typical Use Cases

I can see us mainly wanting to look at a few things to do high angular resolution imaging:

• binary stars: 2x point sources, each with a different spectrum
  • sub-case of this is N point sources for higher order system, each with a different spectrum
  • with wide-ish binaries, field distortion absolutely matters
• star with disk / quasar with background galaxy: 1x point source, 1x pixel grid
  • field distortion will sometimes matter, but not a lot
  • spectra of star/quasar and background resolved source will be very different
• resolved sources with no central star (which might be shrouded in dust): 1x pixel grid
  • field distortion irrelevant (degenerate problem)

We will also sometimes want to do pure engineering work, eg phase retrieval.

Distortion

The best way to think about distortion is that in general a world coordinate system provides an invertible transformation between pixel x, y and world RA, Dec. This may be affine (offset for translation plus pixel scale overall determinant, and rotation, shear) or might be eg polynomial. Aesthetically, I think RA, Dec stuff lives in whatever layer stores information about the astrophysical world - so probably Scene - while the OpticalSystem gets you photons on pixels in x, y coordinates. There is of course the small issue that the geometric optics effect of aberrations is to introduce distortions, so we are separating out the physical optics and geometric optics this way, but ∂Lux isn't a raytracing package (yet?) so this is unavoidable for now. Anyway, we may indeed want an intermediate layer for a Detector!

Regularization

The general way you do image reconstruction is to have a loss function

objective(pixels) = goodness_of_fit(pixels) + alpha*regularizer(pixels)

where alpha is a free, constant, parameter. (You may have a sum of multiple regularizers with multiple coefficients, too!). It is actually quite hard to know in advance what alpha should be - even by orders of magnitude - and there are various hand and adaptive schemes for tuning this. But this is something that happens at the optimization level, and I don't think it necessarily is a property of a layer - is it possible to address specific layers in the loss function? We would like to have nice syntax for doing this.

[LouisDesdoigts]

Distortion is mostly a done-deal in the sense that we shouldn't need to actually develop anything for it (beyond ye-olde' import jax.numpy as np followed by a whack-a-mole of jax arrays do not support update indexing)

Webbpsf actually handles most of the set up (which we can do in the __init__ step) in this script. We might even be able to port a Jax-version of astropy Polynomial2D as an equinox module and use it as an OpticalSystem property

[Jordan-Dennis]

@benjaminpope, with regards to specifying a specific layer in the loss_function the jax.tree_map function will serve you well. This is already how we specify what leaves of the PyTree we want to take the gradient with respect to. Create a clone and then add the alpha parameters onto the leaves that you want to change using the eqx.tree_at syntax.

[ataras2]

I agree with most of what has been said but will also throw in my two cents.

1. I don't mind the name Scene but SkyScene/SkyPatch or similar might work?
2. I think Scene shouldn't call OpticalSystem since this creates the question of ownership and would make it harder to change down the track - which class has the role of making the system? Instead, I think something like this could work, with the OpticalSystem class managing the handover of data and calls to each of Scene, Optics and Detector:
   
   (The reference to Jordan's list is Development of optical components layer script #59)
3. The above also helps with the issue of where to have regularisation (which isn't intuitively a property of Scene) - it should be in the overall OpticalSystem, or perhaps even in a class over the top of it.
4. I quite like the idea of having a (perhaps small) catalogue of scenes as part of dLux - there should be ways to save and load instances of this class easily.

[Jordan-Dennis]

Nice diagram.

[LouisDesdoigts]

Yeah I think we still want to have the OpticalSystem as the main object we interact with, or possibly a create a new class Telescope that contains an OpticalSystem, SkyScene and Detector. It would also make sense to put our distortion code at this level in order to handle the transformations from the Scene to through the optics onto the detector.

This would allow us to basically keep OpticalSystem as it is, letting the more complex programmatic flow needed to model resolved objects be contained in the Telescope class.

This would basically be the same structure that Adam made, but with OpticalSystem/Optics shifted down and replaced with Telescope at the top and with detector able to be inherited by OpticalSystem.

This would also put us more in line with the overall architecture used for webbpsf which is probably not a bad idea in the long run!

[Jordan-Dennis]

Hi all,
I have just started work on the Source objects basing the new branch of the HexagonsAndHexikes (see #67). I thought I would post my current plan here so that @benjaminpope could correct anything he thought needed fixing early on. I have used the new syntax from Patrixk-Kidgers new package (jaxtyping)[https://github.com/google/jaxtyping].

(abstract) class Source(eqx.Module):
    spectrum : f64[2, n] # n is the number of wavelengths in the discrete spectrum.
    position : f64[2, 1] # the x, y position. Silly question but should there be a z?
    resolved : bool # answers the question: is this source resolved?
    
    def __init__(self, spectrum : f64[1, n], position : f4[2, 1]):

    (abstract) def _regularize(self, pixels): # Not sure what is needed to calculate the 
        # regularisation. @benjaminpope, I figure you will have a better idea where this 
        # should and what it will need. 
    
    (abstract) def __call__(self, params_dict): # I'm in favour of generating the 
        # wave-fronts at the source, @LouisDesdoigts may have something to say aboutthis 
        # though. (i.e. replacing CreateWavefront)

class Point(Source):
    def _regularise(self, pixels):
        # call specific @classmethod regulariser. 
    def __call__(self, parameters):

class Smudge(Source): # This is an Extended source Smudge is a joke
    def _regularise(self, pixels):
    def __call__(self, parameters):

From @LouisDesdoigts original post I see there are many ways that the regularisation's can be calculated. I was thinking that these method___s___ could all be implemented in the (abstract) Source class as @classmethods, I would need to check how this interfaced with equinox but it should be fine.

For anyone who may not know, a @classmethod belongs to the class not the objects. In other words it does not depend on the state of the object (you cannot access instances using self.position). From my understanding the regularisation methods do not need to know about the state of the class but rather the psf itself.

The reason I think it belongs in the Source class is because it seems we want to be able to choose the regularisation based on what type of Source we are looking at. Another contending option I can see is that the regularisation methods are split into their own file and structure. e.g.

(abstract) class Regulariser(eqx.Module):
class MaximumEntropy(eqx.Module):
    def __init__(self, ):
    def __call__(self, ):

ect. I do not know enough about regularisation to comment on any shared behaviour or useful functionality at this point, but if someone could comment a list of useful sources I am happy to implement a few of these, provided implementations do not already exist. This seems to tie in with @benjaminpope grand scheme of mixing neural network layers with optical ones willy-nilly.

I'll be messaging again once when I start working on the SkyScene.

Regards

Jordan

[LouisDesdoigts]

Haha I'm still highly amenable to maintaining a general dictionary that is passed between layers, but I'm open to new ideas too if theres good reason.

Yes I imagine compound point source will be a crucial class. One other thing I would note is that we should be somehwat mindful of where the values were trying to optimise will live. I can say from first hand experience you don't want to have to write

filter_spec = eqx.tree_at(lambda osys: osys.sources.source_list[0].binary_mean, osys, True)

every time.

Will just also throw a last plug about an abstract base OpticalSystem, so we don't store position values in two places ✌️

[Jordan-Dennis]

In this case we probably do not need to implement a __call__ method and the sources will just store related information. It also sounds like there are many parameterizations to consider for the Source objects.

I do not have a good way to deal with this in mind. Will Binaries also be parametrised in the same way or do we need to have flexibility? The simplest way is just to have many different implementations of Binary based on the parametrisation.

It should also be possible to translate between the different parametrisations using methods, e.g. sep_dist_and_midpoint_to_absolute, but this could be cumbersome. Additionally you would be learning on the fixed parametrisation and translating the results.

[Jordan-Dennis]

Also @LouisDesdoigts what were you thoughts on the Regularisers

[ataras2]

Re regularisers: I'm strongly in favour of the regulaisers as a seperate class for two main reasons

1. You might want to run autodiff with a particular regularisers for a particular source, but someone else might like to experiment. Keeping the two modular makes this easier
2. I don't think regularisation has to be on just the pixels of the final psf. Could users also add terms in the loss expression to punish unphysical layers? e.g. limiting the relative values in a chromatic phase plate or enforcing sparisty in a layer's contents

[Jordan-Dennis]

Ok that sounds good then. I imagine something like this is widely implemented it just comes down to whether it fits our structure. Do you know if optax has any implemented that we can use, or should we be looking at porting some ourselves?

[benjaminpope]

Yeah I don't think Regularize is necessarily going to be something that lives inside a Source, because there are a few regularizers we could use:

• MaxEnt - sum(p_i * ln(p_i))
• TV - sum over absolute value of finite differences
• Smoothness - sum over square of finite differences
• (neural network learned regularizers in future)

These might be used in combination, too, as a weighted sum!

so I think we need to be able to access the pixel values but not necessarily that these regularizers be built in - or alternately, that the sources by default have multiple regularizer methods.

And as Adam says, we may want to also access other parameters for design problems, where we regularize in various ways phase values etc.

[Jordan-Dennis]

They'll probably get their own file then. Not sure if they will be eqx.Modules or just functions. Probably Modules so that they interface with everything else easily.

[LouisDesdoigts]

Yeah I think it makes the most sense to have the regularisers in their own script as pure functions, kinda of similar to the bayes.py script in that its designed specifically to work with ∂Lux but is just a general jax implementation of some mathematical functionality. This should then let us create more general classes for Source and Scene for simplicity and flexibility.

[ataras2]

At first I would agree that pure functions is the way to go but it also depends what exactly the neural network learned regularizers "learn". Do they learn to approximate a regularisation function for a single layer (or the psf output) or do they learn which layers in an optical system to pay attention to as well as the signals in each layer? I think having them as eqx.Modules might be the way to go if learned regularisers are on the cards.

[Jordan-Dennis]

I'm cautious to using pure functions in the jax ecosystem as they are not valid jax types. This just means that they cannot be passed as first class functions and need to be invoked from within (very hippy 😆).

This means we cannot add regularisers as layers in future, but this is fine if they are only going to be part of larger things.

[Jordan-Dennis]

Hi all,
I thought I would update the discussion by asserting that I had completed an attempt at Source, which I expect will need to be revised by either @LouisDesdoigts or @benjaminpope. I have kept it very simple and each source is little more that a dictionary of its parameters with get_ methods. This was a PEP8 style choice and I am completely amiable to removing them if we decide that it is extra space for this simple class. I also implemented an abstract class (PEP8 again) but also recognise that given the simplicity of the situation this might be overkill and I am not opposed to removing them.

The code is once the Source branch. I have not started work on the SkyScene yet, and would consider it unwise to do so until the Source and Telescope classes are stable, because it will govern interactions between the two.

Please reply to this comment with suggested changes, fixes, additions ect. Please note there was some incomplete code regarding the weights for the combined spectra of binaries. I wasn't sure if we cared about this functionality and the implementation is non-trivial.

[LouisDesdoigts]

Just checked out the code on the Source branch - very happy with how its shaping up. A few things we will want to add too. Firstly we want to some form of spectral parameterisation, ie linear (only a slope), polynomial (some arbitrary number of degrees maybe?), independent (ie free floating individual weights values, not sure about the name). This could possibly be categorised simply as 'polynomial of arbitrary degree' (I do know how you love solving the general case @Jordan-Dennis 😂) and free-floating/independent.

There are also some other functionalities we will want, ie normalisations, possibly regularisations (supporting the idea that we might want regularisations.py as a general package-wide script). We likely also want the option to be able to interpolate the spectrum onto the simulated wavelengths for the filters, or to use degraded/high resolution spectrums interpolating the filter onto those wavelength values for bright objects/ones diffracted to large angles (we have all seen the many lambda/D that those foreground JWST stars diffraction spikes span!).

I suspect this may require a second layer of abstraction for Source to get this all working. Ie the base level source is more-or-less what Source is now, everything array based, with spectral parameterisations handled in the higher level classes, and spectral interpolations handled in the low-level class, which is ultimately what interfaces with Scene/OpticalSystem.

This is a challenging task as there are many different pieces that inform each other, but this looks like were absolutely on the right path here. @benjaminpope maybe you can help @Jordan-Dennis prioritise which parameterisations we want etc. All up some awesome stuff though great work Jordan.

[Jordan-Dennis]

Thanks @LouisDesdoigts, I'll have a crack at the different parametrisations. The hardest part will be naming them all 😆. Did you want me to finish implementing the combined weights for the Binary class?

Also, I don't think it is a problem for now but as with the propagators we run the risk of having a very large number of classes. i.e. LinearAndQuadraticBinary, QuadraticBinary, LinearBinary, LinearAndArbitraryBinary ect. While this isn't necessarily a bad thing I wonder if we shouldn't try and roll some of the functionaloty together somehow.

[LouisDesdoigts]

I think we should be able to do all the spectral parameterisation from the base source class, and then specify which type to use for each individual in the constructor, with each sources' spectrum treated independently. Each source from a binary star has an independent spectrum, so we shouldn't need any of those classes I believe. We only need to parameterise their positions since they are gravitationally bound objects.

I've also been thinking some more, I believe we only need to implement one general case nth-degree polynomial spectrum and an array spectrum (no parametrisation), linear is just a degree-1 polynomial lol.

[Jordan-Dennis]

That sounds good, I think we are on the same page now. I assume Taylor approximations will be fine for things like exp and log if we want to model spectra like that or do those situations simply not arise?

[LouisDesdoigts]

I don't think there are any cases we would want more anything more complex than a polynomial. We can always add it later if we need, but definitely not worth putting any effort into it now.

[LouisDesdoigts]

Just throwing a small plug here for #70 - I think this is something that would be useful, and could inform/help guide how we build the Source objects since it may not be practical to fully construct a full scene object at once, who knows though!

[Jordan-Dennis]

Hi Guys,
Just thought I would revive this discussion. So I am not sure how we are planning to interface between the Sources and the Wavefronts. For example, if I have a two Sources parameterised by their angular coordinates, how does this get translated into a pair of Wavefronts? Assuming they are both equally, infinitely distant, we will get a constant amplitude and 0 relative phase for both. This implies that they will be propagated by the telescope in the same way and as a result the PSF will be in exactly the same. I imagine that the difference comes in how the coordinates underlying the transformations are generated. This complicates matters because it implies that the Source needs to know/cause the state of the Wavefront; specifically the get_pixel_positions. This also adds the possibility of additional parameters in Wavefront keeping track of the paraxial_centre and brightness that will be required to correctly super-impose the output Wavefronts at the Detector.

There is also the possibility that the a single PSF is generated based on the optical transmission of the Telescope and that this is then (brightness) scaled and translated according to the information in Source. I imagine this is probably the way we will go. If so the question remains how do we determine the position of the PSF in the Detector image based on the coordinates of the Source; particularly, does this depend on the Telescope implementation. i.e. does the position of the PSF change based on whether the image was taken using JamesWebbSpaceTelescope or HubbleSpaceTelescope.

As always let me know your thoughts.

Jordan

[benjaminpope]

Sio if you have 2 sources with different positions, the positions are encoded by a phase slope across their wavefront.

[Jordan-Dennis]

Very small slope given the distances I imagine. Based on this I think that we need to consider replacing CreateWavefront with Source and add a set_units/set_type method to Scene. I think that it will be inefficient and complicated to have Sources pass information to CreateWavefront which will also need to know the Wavefront type via an additional parameter of some sought.

@LouisDesdoigts what do you think?

[benjaminpope]

Well 1 λ slope across the pupil will shift the PSF centre by λ/D. So it's not actually that small in phase terms but it's small in nm.

(This is how we position the binaries already in the fisher information example.

[Jordan-Dennis]

Ok cool, I will draft and example of this over the coming days (maybe Sunday). The present Source code is very rudimentary and pretty much just lists the parameters.

[LouisDesdoigts]

Yeah so just to jump back in here, I don't think the issues you've brought up @Jordan-Dennis will be an issue for us.

Firstly as @benjaminpope mentioned, all point source positions can only be parameterised by their offset angle from the pointing axis of the telescope, ie relative to the optical axis. Every point within an optical system along the optical axis has some linear relationship between its input offset angle and the corresponding angle that it is deflected away from that axis. Ie a high field of view optical system will 'deflect' some input source a smaller angle than a small field of view optical system. That angle is determined by the properties of all the preceding optical elements, but no matter the optical system that shift is exclusively parametrised by the phase slope across the aperture, and hence the angular offset from the optical axis.

This is great for us because it means that source and scene are entirely agnostic to the optical system observing it. Source and scene can store and parametrise on-sky properties arbitrarily and only need to spit out a series of angular positions, which can then be turned into angular positions relative to the optical axis defined by the telescope pointing. This is all you need to create and propagate the wavefront, allowing source/scene able to be entirely ignorant of the optical system.

A similar logic applies to optical system and detector. Optical system output an array of flux values entirely parameterised by the pixel scale since all psfs are produced relative to the paraxial axis. No need to be aware of anything else going on upstream.

This simply means that the telescope class does the heavy lifting, transforming information between these distinct classes that are agnostic to each other. It takes input angular positions and relative spectral fluxes and individually propagates these through its optical system, and then passes them to the detector. This means it does all the formatting of psfs in terms of relative brightnesses, 'observing strategies', etc.

Let me know if that makes sense, or if I've perhaps misunderstood some nuance your original comment, but source and scene should be relatively simple in my understanding. I think they only need to take a series of arbitrarily parameterised on sky sources and spit out relative Cartesian position angles, and spectral fluxes ✌️