Define your own simulation element

Marxs comes with a selection of useful optical elements and for many cases these elements are sufficient to define your mission. However, if that is not the case, you may want to write your own code to deal with those pieces that marxs is missing. We are very interested in adding more elements to marxs and we would love to hear from you. Just open an issue or pull request to add your code to marxs at our github repro.

In the simplest case, an element of the simulation can be any function or class with a __call__ method that accepts a photon list as input and returns a (usually modified) photon list as output. As a simple example we show an element that represent an absorbing plate that just absorbes half of all photons coming though (i.e. reduces the probability of each photon to be detected by 1/2). First, let us generate some photons coming from some astrophysical source:

>>> from astropy.coordinates import SkyCoord
>>> from marxs import source, optics
>>> import astropy.units as u
>>> mysource = source.PointSource(coords=SkyCoord(181., -23., unit='deg'))
>>> photons = mysource.generate_photons(10 * u.s)
>>> mypointing = source.FixedPointing(coords=SkyCoord('12h01m23s -22d59m12s'))
>>> photons = mypointing(photons)

Our example is very simple. We can do this with a simple function:

>>> def myabs1(photons):
...     photons['probability'] *= 0.5
...     return photons
>>> photons = myabs1(photons)

A more complicated way to achieve the same thing would be:

>>> from marxs.optics import OpticalElement
>>> class AbsorbAllPhotonsElement(OpticalElement):
...     def __call__(self, photons):
...         photons['probability'] *= 0.5
...         return photons
>>> myabs2 = AbsorbAllPhotonsElement()
>>> photons = myabs2(photons)

In this simple case, a function might be sufficient but most optical elements are more complex. For example, they have a finite size, they cannot be vectorized and must loop over all photons one-by-one, or they should offer the user some diagnostic to analyze the result of the ray-trace such as an index column that lists which CCD was hit or which mirror shell reflected a particular photon. For those cases, marxs provides base classes that do much of the work:

Check the description of the classes in this list for details. Optical elements in marxs are derived from OpticalElement or FlatOpticalElement. The OpticalElement requires the user to implement an intersect method (see intersect for a describtion of input and output parameters). Then either of the following will work:

  • Override process_photons with the method signature.

  • For a vectorized approach, just add a method called specific_process_photons, that has the same input signature as process_photons and outputs a dictionary, where the keys are columns names in the photon list and the values are numpy.ndarray objects that hold the result for the intersecting photons.

  • Instead, if the derived element has a method called process_photon (see process_photon for the required signature) then it will loop through the list of intersecting photons one-by-one and record the output values of that function in the photons table.

A FlatOpticalElement already has an intersect method implementing the default box geometry discussed in Coordinate system, so the best us is to implement a specific_process_photons method.

For performance, we recommend to use vecorized implementation and change the values of the photon table in place (as opposed to making copies of parts of the table and joining them back together) whereever possible.

Example for a derived class

A simple example for an optical element is the baffle - a rectangular opening in a large metal plate. It can be inserted into the beam and all photons that do not pass through the hole will be absorbed. In practice, the metal plate is not infinite, but has an outer bound, too, but all photons passing on the outside are either absorbed by the satellite housing ot the tube of the lab beamline and are lost from the experiment anyway and do not need to be traced any further.

# Licensed under GPL version 3 - see LICENSE.rst
import numpy as np

from .base import FlatOpticalElement
from ..math.utils import h2e
from ..visualization.utils import plane_with_hole


class Baffle(FlatOpticalElement):
    '''Plate with rectangular hole that allows photons through.

    The probability of photons that miss is set to 0.

    Parameters
    ----------
    photons: astropy Table
        table that includes information on all of the photons
    '''

    display = {'color': (1., 0.5, 0.4),
               'outer_factor': 3,
               'shape': 'plane with hole'}


    def process_photons(self, photons, intersect, intercoos, interpoos):
        photons['pos'][intersect] = intercoos[intersect]
        photons['probability'][~intersect] = 0
        return photons

    def triangulate_inner_outer(self):
        '''Return a triangulation of the baffle hole embedded in a square.

        The size of the outer square is determined by the ``'outer_factor'`` element
        in ``self.display``.

        Returns
        -------
        xyz : np.array
            Numpy array of vertex positions in Eukeldian space
        triangles : np.array
            Array of index numbers that define triangles
        '''
        r_out = self.display.get('outer_factor', 3)
        g = self.geometry
        outer = h2e(g['center']) + r_out * np.vstack([h2e( g['v_x']) + h2e(g['v_y']),
                                                      h2e(-g['v_x']) + h2e(g['v_y']),
                                                      h2e(-g['v_x']) - h2e(g['v_y']),
                                                      h2e( g['v_x']) - h2e(g['v_y'])
        ])
        inner =  h2e(g['center']) + np.vstack([h2e( g['v_x']) + h2e(g['v_y']),
                                               h2e(-g['v_x']) + h2e(g['v_y']),
                                               h2e(-g['v_x']) - h2e(g['v_y']),
                                               h2e( g['v_x']) - h2e(g['v_y'])
        ])
        return plane_with_hole(outer, inner)

The code above uses the intersect method it inherits to calculate which photons intersect at all and if so, at which coordinate. This point is entered into the photons['pos'].

Note that the baffle can be rectangular without writing a single line of code here. baf = Baffle(zoom=[1,3,2]) would create a baffle of size 6*4 (the outer boundaries are 3 or 2 mm from the center) because the OpticalElement can interpret all Coordinate system keywords and the marxs.optics.FlatOpticalElement.intersect makes use of this information when it calculates intersection points.

The baffle does not add any new descriptive columns to the output, thus the value id_col defined in a base class is not overwritten. However, the baffle does have a class dictionary display that contains default parameters for plotting (see Visualize results).

Non-physical elements

Above, we discussed elements that are physically present, such as mirrors or detectors. The simulations also allows elements that are not a piece of hardware, for example, a function that just saves a copy of the photon list to disk or prints some diagnostic information would also work:

>>> def backupcopy(photons):
...     photons.write('backup.fits', overwrite=True)
...     return photons
>>> photons = backupcopy(photons)

Why have we chosen to return a copy of photons here, although no value in the photon list was changed? This allows us to sneak in the new function into a list of optical elements:

>>> from marxs import simulator
>>> myinstrument = simulator.Sequence(elements=[myabs1, backupcopy, myabs2])

Reference / API

marxs.base Package

Functions

check_energy_consistent(photons)

Check if the energy is the same for all photons.

check_meta_consistent(meta1, meta2[, ...])

Check that the meta data between two simulations indicates consistency.

Classes

DocMeta(name, bases, dict)

Metaclass to inherit docstrings when required.

GeometryError

MarxsElement(**kwargs)

Base class for all elements in a MARXS simulation.

SimulationSequenceElement(**kwargs)

Base class for all elements in a simulation that processes photons.

TagVersion([ORIGIN, CREATOR])

Tag a photons list with diagnostic information such as a the program version.

Class Inheritance Diagram

Inheritance diagram of marxs.base.base.DocMeta, marxs.base.base.GeometryError, marxs.base.base.MarxsElement, marxs.base.base.SimulationSequenceElement, marxs.base.base.TagVersion