# Working with dLux Objects

This tutorial is designed to give users a quick overview of how to work with dLux objects. Build using `Equinox`

and `Zodiax`

, dLux objects intuitive and simple to work with, so lets have a look at how to get started.

```
# Basic imports
import jax.numpy as np
import jax.random as jr
# dLux imports
import dLux as dl
import dLux.utils as dlu
```

## An Optical System

First we set up a dLux object to work with, in this case a simple optical system. We will not cover the details of how to build an optical system here, as it is covered elsewhere in the tutorials.

```
# Define our wavefront properties
wf_npix = 512 # Number of pixels in the wavefront
diameter = 1.0 # Diameter of the wavefront, meters
# Construct a simple circular aperture
coords = dlu.pixel_coords(wf_npix, diameter)
aperture = dlu.circle(coords, 0.5 * diameter)
# Zernike aberrations
indices = np.array([2, 3, 7, 8, 9, 10])
basis = 1e-9 * dlu.zernike_basis(indices, coords, diameter=diameter)
coefficients = 50 * jr.normal(jr.PRNGKey(0), indices.shape)
# Define our detector properties
psf_npix = 64 # Number of pixels in the PSF
psf_pixel_scale = 50e-3 # 50 mili-arcseconds
oversample = 3 # Oversampling factor for the PSF
# Define the optical layers
# Note here we can pass in a tuple of (key, layer) paris to be able to
# access the layer from the optics object with the key!
layers = [
(
"aperture",
dl.layers.BasisOptic(
transmission=aperture,
basis=basis,
coefficients=coefficients,
normalise=True,
),
),
dl.layers.Tilt(np.zeros(2)),
]
# Construct the optics object
optics = dl.AngularOpticalSystem(
wf_npix, diameter, layers, psf_npix, psf_pixel_scale, oversample
)
# Let examine the optics object! The dLux framework has in-built
# pretty-printing, so we can just print the object to see what it contains.
print(optics)
```

```
AngularOpticalSystem(
wf_npixels=512,
diameter=1.0,
layers={
'aperture':
BasisOptic(
basis=f32[6,512,512],
coefficients=f32[6],
as_phase=False,
transmission=f32[512,512],
normalise=True
),
'Tilt':
Tilt(angles=f32[2])
},
psf_npixels=64,
oversample=3,
psf_pixel_scale=0.05
)
```

## Paths

So now that we have our optical system set up and we can see the layout, lets have a look at how to work with them. Being built in `Zodiax`

, dLux gains access to a 'path-based' interface, greatly simplifying how we work with these objects.

A path in `Zodiax`

works very similarly to a path in a file system. It is a way of navigating through the object, and accessing the data we want via strings, joined with a dot ('.'). Here are some example paths for our optical system:

`'diameter'`

`'layers.Tilt.angles'`

`'layers.aperture.transmission'`

dLux also makes extensive use of the `__getattr__`

methods, which allows for the raising of low-level attributes to the top level object. Primarily this means we can skip the `'layers'`

part of these paths, so the above paths become:

`'diameter'`

`'Tilt.angles'`

`'aperture.transmission'`

Now to access these values at these paths, we can either use the `.get(path)`

method, or just access the via the regular attribute accessors. Lets have a look at this in practice:

```
# Using the regular accessors
print("Diameter: ", optics.diameter)
print("Angles: ", optics.Tilt.angles)
print("Transmission", optics.aperture.transmission)
```

```
Diameter: 1.0
Angles: [0. 0.]
Transmission [[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
...
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]]
```

```
# Using the .get method
print("Diameter: ", optics.get('diameter'))
print("Angles: ", optics.get('Tilt.angles'))
print("Transmission", optics.get('aperture.transmission'))
```

```
Diameter: 1.0
Angles: [0. 0.]
Transmission [[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
...
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]]
```

So, if we can access the attributes with regular accessors, what is the point of the `Zodiax`

`.get`

method? Well the `.get`

method lets us access *multiple* attributes at once by passing in a *list* of paths. Lets have a look at this in practice:

```
print(optics.get(['diameter', 'Tilt.angles', 'aperture.transmission']))
```

```
[1.0, Array([0., 0.], dtype=float32), Array([[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]], dtype=float32)]
```

These paths can actually be simplified *further*, since dLux also raises attributes from the *values* of the layers dictionary, so the above paths become:

`'diameter'`

`'transmission'`

`'angles'`

Lets look at this in practice.

**NOTE**

While this level parameter raising can greatly simplify the paths we work with, we need to cognisant that each path is *unique*. For example if we have two layers that have the same values `as_phase`

, then using `'as_phase'`

as our path will only return *one* of these values. To distinguish between these two we would need to reference the layer by its dictionary key, ie `'layer1.as_phase'`

or `'layer2.as_phase'`

.

```
print(optics.get(['diameter', 'angles', 'transmission']))
```

```
[1.0, Array([0., 0.], dtype=float32), Array([[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
...,
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]], dtype=float32)]
```

## Zodiax Methods

`Zodiax`

gives us access to a series of class methods that are designed to mirror the `jax.Array`

syntax, ie `.at`

, `.set`

, `.multiply`

etc. The syntax a slightly different in that we need to specify a path to the data we want to work with, but the functionality is the same. Here are the main `Zodiax`

methods:
some of the zodiax methods:

`.get(paths)`

`.set(paths, values)`

`.add(paths, values)`

`.multiply(paths, values)`

`.divide(paths, values)`

`.min(paths, values)`

`.max(paths, values)`

Lets use the `.add`

method to see how we can modify our optical system.

```
paths = ['diameter', 'angles']
new_optics = optics.add(paths, [1, 0.5])
print(new_optics.get(paths))
```

```
[2.0, Array([0.5, 0.5], dtype=float32)]
```

## Nesting

Zodiax goes further here, and allows to 'nest' paths, such that we can operate on *mulitple* values in the same operation. Lets look at some examples of this in practice:

```
# Operate on multiple values simultaneously
paths = ['diameter', 'angles']
new_optics = optics.multiply(paths, 0)
print(new_optics.get(paths))
```

```
[0.0, Array([0., 0.], dtype=float32)]
```

We can also nest *within* our path itself, lets see how:

```
# Set nested values simulatenously
# Note that 'paths' here has two entries, so we need to supply a
# list of values of the same length
paths = [['diameter', 'angles'], 'transmission']
values = [1, 2]
new_optics = optics.set(paths, values)
print(new_optics.get(paths))
```

```
[1, 1, 2]
```

Excellent! Now some keen readers may have noticed that using the `.set`

method, the `transmission`

values has changes from an array to a float! This is becuase there is no *robust* way to do runtime type and shape checking, plus sometimes we may want to change the type anyway. Do be cognisant when setting values that you are setting then to valid types for the object.

**Summary**

So that is a quick overview of how to work with dLux objects. We have seen how to access the data via paths, and how to modify the data using the `Zodiax`

methods. We have also seen how to nest paths, and how to nest paths within paths. Hopefully this has given you a good overview of how to work with dLux objects, and you can now go and work with the objects with ease!

TODO: Add jit to this tutorial?

TODO: Add object vectorsation?