diff --git a/CHANGES.rst b/CHANGES.rst index 195b60c..c2bac2d 100644 --- a/CHANGES.rst +++ b/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 ^^^^^^^^^^^ diff --git a/marxs/design/tests/test_tolerancing.py b/marxs/design/tests/test_tolerancing.py index e5cfbac..d2b524e 100644 --- a/marxs/design/tests/test_tolerancing.py +++ b/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. diff --git a/marxs/design/tolerancing.py b/marxs/design/tolerancing.py index 74a2409..696dc44 100644 --- a/marxs/design/tolerancing.py +++ b/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. diff --git a/marxs/optics/detector.py b/marxs/optics/detector.py index 148a2c1..c536bf2 100644 --- a/marxs/optics/detector.py +++ b/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 diff --git a/marxs/simulator/simulator.py b/marxs/simulator/simulator.py index 37f8627..55861f3 100644 --- a/marxs/simulator/simulator.py +++ b/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. diff --git a/marxs/simulator/tests/test_simulator.py b/marxs/simulator/tests/test_simulator.py index 18d6946..9bce153 100644 --- a/marxs/simulator/tests/test_simulator.py +++ b/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)