Tolerancing

CHANGES.rst

@@ -4,6 +4,12 @@
 New Features
 ^^^^^^^^^^^^
 
+- `marxs.simulator.Parallel` and `marxs.simulator.Sequence` now have a
+  ``first_of_class_top_level`` method to help find certain subelements. [#222]
+
+- Added `marxs.design.tolerancing.run_tolerances_for_energies2` is similar to
+  `marxs.design.tolerancing.run_tolerances_for_energies`, but with a
+  simplified interface. [#222]
 
 API Changes
 ^^^^^^^^^^^

marxs/design/tests/test_tolerancing.py

@@ -9,13 +9,14 @@
 from astropy.utils.data import get_pkg_data_filename
 
 from ..tolerancing import (oneormoreelements,
-                           wiggle, moveglobal, moveindividual,
+                           wiggle, moveglobal, moveindividual, moveelem,
                            varyperiod, varyorderselector, varyattribute,
                            run_tolerances, CaptureResAeff,
                            generate_6d_wigglelist,
                            select_1dof_changed,
                            plot_wiggle, load_and_plot,
                            run_tolerances_for_energies,
+                           run_tolerances_for_energies2,
                            )
 from ...optics import (FlatGrating, OrderSelector, RadialMirrorScatter,
                        RectangleAperture, ThinLens, FlatDetector)
@@ -114,6 +115,15 @@ def test_moveelements_rotate():
                            np.stack([e.pos4d for e in g2.elements]))
 
 
+def test_moveelem():
+    '''Move an individual element'''
+    det = FlatDetector(zoom=[1, 100, 100])
+    assert det.geometry['center'][2] == 0
+    moveelem(det, dz=5)
+    assert det.geometry['center'][2] == 5
+    assert np.all(det.geometry.pos4d[:3, :3] == np.eye(3) * [1, 100, 100])
+
+
 def test_wiggle():
     '''Check wiggle function'''
     g = gsa()
@@ -213,10 +223,11 @@ def afunc(photons):
     assert 'meanorder' not in parameters[0]
 
 def test_run_tolerances_for_energies():
-    '''For this test, we need to define an instrument. The instrument is not very
-    realistic (an X-ray mirror with r=0 won't work), but the point here is just to
-    check the tolerancing for several energies. To make that calculation reasonably
-    fast, we need to keep the number of elements in the optical system small.
+    '''For this test, we need to define an instrument. The instrument is not
+    very realistic (an X-ray mirror with r=0 won't work), but the
+    point here is just to check the tolerancing for several
+    energies. To make that calculation reasonably fast, we need to
+    keep the number of elements in the optical system small.
     '''
     coords = SkyCoord(12. * u.deg, -45 * u.deg)
     src = PointSource(coords=coords)
@@ -251,6 +262,43 @@ def test_run_tolerances_for_energies():
     assert not np.any(np.isfinite(res['R'].data[:, 2]))
     assert res['R'].data[2, 1] > res['R'].data[0, 1]
 
+
+def test_run_tolerances_for_energies2():
+    '''Same, as above, but with different calling sequence
+    '''
+    coords = SkyCoord(12. * u.deg, -45 * u.deg)
+    src = PointSource(coords=coords)
+    pnt = FixedPointing(coords=coords)
+    aper = RectangleAperture(position=[5000, 0, 0], zoom=[1, 10, 10])
+    lens = ThinLens(position=[4900, 0, 0], zoom=[1, 10, 10], focallength=4900)
+    grat = FlatGrating(d=.002, order_selector=OrderSelector([0, 1]),
+                       position=[4800, 0, 0], zoom=[1, 10, 10])
+    det = FlatDetector(zoom=[1, 100, 100])
+    instrum = Sequence(elements=[pnt, aper, lens, grat, det])
+
+    parameters = [{'period_mean': 0.003, 'period_sigma': 0.},
+                  {'period_mean': 0.004, 'period_sigma': 0.}]
+
+    res = run_tolerances_for_energies2(src, [.1, 1] * u.keV,
+                                       instrum, FlatGrating,
+                                       varyperiod,
+                                       parameters,
+                                       CaptureResAeff(orders=[0, 1, 2]),
+                                       reset={'period_mean': 0.005,
+                                              'period_sigma': 0.},
+                                       t_source=1. * u.ks)
+    # Check the reset worked
+    assert grat._d == 0.005
+    # Check both energy have been calculated
+    assert 1 in res['energy']
+    assert .1 in res['energy']
+    assert len(res) == 4
+    # check results are reasonable
+    assert np.all(res['R'].data[:, 0] == 0)
+    assert not np.any(np.isfinite(res['R'].data[:, 2]))
+    assert res['R'].data[2, 1] > res['R'].data[0, 1]
+
+
 def test_capture_res_aeff():
     '''Test the captures res/aeff class.
 

marxs/design/tolerancing.py

@@ -10,14 +10,14 @@
 from astropy.table import Table
 from astropy import table
 import astropy.units as u
-from ..simulator import Parallel
+from ..simulator import Parallel, Sequence
 from .uncertainties import generate_facet_uncertainty as genfacun
 from ..analysis.gratings import (resolvingpower_from_photonlist,
                                  effectivearea_from_photonlist)
 from ..analysis.gratings import AnalysisError
 
 __all__ = ['oneormoreelements',
-           'wiggle', 'moveglobal', 'moveindividual',
+           'wiggle', 'moveglobal', 'moveindividual', 'moveelem',
            'varyperiod', 'varyorderselector', 'varyattribute',
            'run_tolerances', 'run_tolerances_for_energies',
            'CaptureResAeff',
@@ -105,6 +105,36 @@ def moveindividual(e, dx=0, dy=0, dz=0, rx=0, ry=0, rz=0):
     e.generate_elements()
 
 
+
+@oneormoreelements
+def moveelem(e, dx=0, dy=0, dz=0, rx=0., ry=0., rz=0.):
+    '''Move and rotate a marxs element around principal axes.
+
+    Unlike `~marxs.design.tolerancing.moveglobal` this does not work on a
+    `~marxs.simulator.Parallel` object, which is a container holding other
+    elements, but it moves just one specific element itself (e.g. a single
+    detector).
+
+    To make it easier to return the element to its original position,
+    the original ``pos4d`` matrix of the element is stored as ``pos4d_orig``.
+
+    Parameters
+    ----------
+    e :`marxs.optics.base.OpticalElement` or list of those elements
+        Elements where uncertainties will be set
+    dx, dy, dz : float
+        translation in x, y, z (in mm)
+    rx, ry, rz : float
+        Rotation around x, y, z (in rad)
+    '''
+    if not hasattr(e.geometry, 'pos4d_orig'):
+        e.geometry.pos4d_orig = e.geometry.pos4d.copy()
+    move = affines.compose([dx, dy, dz],
+                           euler.euler2mat(rx, ry, rz, 'sxyz'),
+                           np.ones(3))
+    e.geometry.pos4d = move @ e.geometry.pos4d_orig
+
+
 @oneormoreelements
 def varyattribute(element, **kwargs):
     '''Modify the attributes of an element.
@@ -345,7 +375,7 @@ def run_tolerances_for_energies(source, energies,
                                 instrum_before, instrum_remaining,
                                 wigglefunc, wiggleparts,
                                 parameters, analyzefunc,
-                                reset=None, t_source=1):
+                                reset=None, t_source=1 / u.s):
     '''Run tolerancing for a grid of parameters and energies
 
     This function loops over `~marxs.design.tolerancing.run_tolerances` for
@@ -397,7 +427,7 @@ def run_tolerances_for_energies(source, energies,
         If ``reset=None``, then the elements affected by ``wigglefunc`` will
         be in the state corresponding to the last entry in ``parameters`` when
         this function exits.
-    t_source : int
+    t_source : `astropy.units.quantity.Quantity`
         parameter for ``source.generate_photons()``. If the source flux is set
         to one, then ``t_source`` determines the number of photons used in the
         tolerancing simulation.
@@ -434,6 +464,82 @@ def run_tolerances_for_energies(source, energies,
     return table.vstack(outtabs)
 
 
+# This function could be moved marxs itself
+def run_tolerances_for_energies2(source, energies, instrum, cls, wigglefunc,
+                                 parameters,
+                                 analyzefunc,reset=None, subclass_ok=False,
+                                 t_source=1 / u.s):
+    '''Run tolerancing for a grid of parameters and energies
+
+    This function loops over `~marxs.design.tolerancing.run_tolerances` for
+    different energies.
+    This function takes an instrument configuration and a function to change
+    one aspect of it. For every change it runs a simulations and calculates a
+    figure of merit. As the name indicates, this function is designed to derive
+    alignment tolerances for certain instrument parts but it might be general
+    enough for other parameter studies in instrument design.
+
+    There are two nested loops here looping over energy and tolerancing
+    parameters. In this case, the outer loop is over energies and then for
+    every energy, `~marxs.design.tolerancing.run_tolerances` loops over the
+    parameters. This works well where running the ``wigglefunc`` is relatively
+    fast, but propagating the photons through ``instrum_before`` is slow and
+    minimizes the memory footprint, because only one photon list is in memory
+    at any one time. It also implies that runs for the same wiggle parameters
+    but different energies are run on different realizations of any random
+    draws that are performed by ``wigglefunc``.
+
+
+    Parameters
+    ----------
+    source : `marxs.source.Source`
+        Source used to generate the photons for every energy. This function
+        changes the energy of the source, so the source should be set for
+        monoenergetic emission with a constant flux of 1.
+    energies : `astropy.units.quantity.Quantity`
+        An array of energy values.
+    instrum : `marxs.simulator.Sequence`
+        An instance of the instrument which contains all elements.
+    cls : class
+        Class of modified / wiggled / etc. element
+    wigglefunc, parameters, analyzefunc :
+        These parameters are passed to
+        `marxs.design.tolerancing.run_tolerances`. See that function for a
+        description of these parameters.
+    reset : dict or ``None``
+        A dictionary of values for the ``wigglefunc`` function that resets
+        the positions or properties of the wiggled elements to their default.
+        If ``reset=None``, then the elements affected by ``wigglefunc`` will
+        be in the state corresponding to the last entry in ``parameters`` when
+        this function exits.
+    t_source : `astropy.units.quantity.Quantity`
+        parameter for ``source.generate_photons()``. If the source flux is set
+        to one, then ``t_source`` determines the number of photons used in the
+        tolerancing simulation.
+
+    Returns
+    -------
+    result : `astropy.table.Table`
+        Each row in the table contains energy, wave, parameter values, and
+        results from ``analyzefunc`` for a single run.
+
+    '''
+    ind = instrum.first_of_class_top_level(cls, subclass_ok)
+    if ind is None:
+        raise Exception(f'{cls} nor part of {instrum}')
+
+    tab = run_tolerances_for_energies(source, energies,
+                                      Sequence(elements=instrum.elements[:ind]),
+                                      Sequence(elements=instrum.elements[ind:]),
+                                      wigglefunc,
+                                      instrum.elements_of_class(cls),
+                                      parameters,
+                                      analyzefunc,
+                                      reset,
+                                      t_source=t_source)
+    return tab
+
+
 def generate_6d_wigglelist(trans, rot,
                            names=['dx', 'dy', 'dz', 'rx', 'ry', 'rz']):
     '''Generate a list of parameters for the wiggle functions in this module.

marxs/optics/detector.py

@@ -1,5 +1,4 @@
 # Licensed under GPL version 3 - see LICENSE.rst
-from __future__ import division
 import warnings
 
 import numpy as np

marxs/simulator/simulator.py

@@ -99,6 +99,48 @@ def elements_of_class(self, cls, subclass_ok=False, stop_at_first=False):
                 a += e.elements_of_class(cls, subclass_ok)
         return a
 
+
+    def first_of_class_top_level(self, cls, subclass_ok=False):
+        '''Return the index of the first sub-element that has ``cls`` in it.
+
+        This walks the hiracy of `marxs.simulator.Sequence` and
+        `marxs.simulator.Parallel` objects and returns the index of the
+        first sub-element that contains a certain class.
+
+        Typically, this function is used to study and instrument with small
+        modifications. If e.g. only the detector are changed, we might want to
+        run a simulation up to the detectors, save te rays and then try out
+        different detector combinations without rerunning the simulations for
+        all earlier steps. This methid can help to identify the step that
+        contains the detectors.
+
+        Parameters
+        ----------
+        cls : class
+            The class searched for, e.g. `marxs.optics.FlatDetector` to find all
+            detectors.
+        subclass_ok : bool
+            Controls if subclasses of `cls` count as match.
+
+        Returns
+        -------
+        out : int or None
+            index for the first element in ``self.elements`` that is or contains
+            an objects of class ``cls``. If none is found, the method return
+            ``None``.
+        '''
+        for i, e in enumerate(self.elements):
+            if subclass_ok:
+                compfunc = lambda a, b: isinstance(a, b)
+            else:
+                compfunc = lambda a, b: type(a) == b
+            if compfunc(e, cls):
+                return i
+            elif hasattr(e, 'elements_of_class') and len(e.elements_of_class(cls, subclass_ok)) > 0:
+                return i
+        return None
+
+
 class Sequence(BaseContainer):
     '''A `Sequence` is a container that summarizes several optical elements.
 

marxs/simulator/tests/test_simulator.py

@@ -3,7 +3,7 @@
 import pytest
 
 from ..simulator import KeepCol, Sequence, BaseContainer
-from ...optics import FlatDetector
+from ...optics import FlatDetector, FlatOpticalElement
 
 f1 = FlatDetector()
 f2 = FlatDetector()
@@ -35,6 +35,14 @@ def test_seach_top():
                                                      stop_at_first=True)
 
 
+def test_first_of_class():
+    '''Check that the first of class method returns expected indices.'''
+    assert mission.first_of_class_top_level(FlatDetector) == 0
+    assert mission.first_of_class_top_level(FlatOpticalElement) is None
+    assert mission.first_of_class_top_level(FlatOpticalElement, subclass_ok=True) == 0
+    assert mission.first_of_class_top_level(Sequence) == 1
+
+
 def test_format_saved_positions():
     '''Reformat saved positions and drop nearly identical values.'''
     pos0 = np.arange(20).reshape(5,4)