# Licensed under GPL version 3 - see LICENSE.rst
'''Gratings and efficiency files'''
import math
import numpy as np
from ..math.utils import norm_vector
from ..math.polarization import parallel_transport
from .. import energy2wave
from .base import FlatOpticalElement
[docs]
class OrderSelector(object):
'''Select from a list of order numbers independent of energy.
Parameters
----------
orderlist : array
This are the order numbers to chose from. They must be integers.
p : None or array
Probability for a photon to end up in each order. The sum of
all probabilities can be smaller than 1, if a certain fraction
of photons is absorbed by the grating. If this is ``None``,
equal probability is assigned to all orders.
Examples
--------
Two common use cases for testing are a grating where every photon gets
diffracted into the same order or where all photons get distributed with
equal probability into a set of orders.
>>> import numpy as np
>>> from marxs.optics import FlatGrating, OrderSelector
>>> singleordergrating = FlatGrating(d=1e-4, order_selector=OrderSelector([3]))
>>> setofordersgrating = FlatGrating(d=1e-4, order_selector=OrderSelector(np.arange(-3, 4)))
'''
def __init__(self, orderlist, p=None):
self.orderlist = orderlist
if p is None:
self.p = np.ones(len(orderlist)) / len(orderlist)
else:
p = np.asanyarray(p)
if len(p) != len(orderlist):
raise ValueError('Number of elements in orderlist and probabilities does not match.')
if p.sum() > 1.:
raise ValueError('Sum of all probabilities must be <= 1.')
if np.any(p < 0):
raise ValueError('Probabilities cannot be negative')
self.p = p
[docs]
def __call__(self, energy, *args):
p_sum = self.p.sum()
p = self.p / p_sum
if np.isscalar(energy):
return np.random.choice(self.orderlist, p=p), self.p.sum()
else:
return np.random.choice(self.orderlist, size=len(energy), p=p), p_sum * np.ones_like(energy)
[docs]
class EfficiencyFile(object):
'''Select grating order from a probability distribution in a data file.
The file format supported by this class is as follows: The first
colum contains energy values in keV, all remaining columns have
the probability that a photons with this energy is diffracted into
the respective order. The probabilities for each order do not have
to add up to 1.
Parameters
----------
filename : string
Path to the efficiency file.
orders : list
List of orders in the file. Must match the number of columns with
probabilities.
'''
def __init__(self, filename, orders):
dat = np.loadtxt(filename)
self.energy = dat[:, 0]
if len(orders) != (dat.shape[1] - 1):
raise ValueError('orders has len={0}, but data files has {1} order columns.'.format(len(orders), dat.shape[1] - 1))
self.orders = np.array(orders)
self.prob = dat[:, 1:]
# Probability to end up in any order
self.totalprob = np.sum(self.prob, axis=1)
# Cumulative probability for orders, normalized to 1.
self.cumprob = np.cumsum(self.prob, axis=1) / self.totalprob[:, None]
[docs]
def __call__(self, energies, *args):
orderind = np.empty(len(energies), dtype=int)
ind = np.empty(len(energies), dtype=int)
for i, e in enumerate(energies):
ind[i] = np.argmin(np.abs(self.energy - e))
orderind[i] = np.min(np.nonzero(self.cumprob[ind[i]] > np.random.rand()))
return self.orders[orderind], self.totalprob[ind]
[docs]
class FlatGrating(FlatOpticalElement):
'''Flat grating
The grating is assumed to be geometrically thin, i.e. all photons enter on
the face of the grating, not through the sides.
Parameters
----------
d : float or callable
grating constant in mm. If ``d`` is callable, then it will be
called as ``d(intercoos)``, where ``intercoos`` is a (N, 2)
array holding the positions where photons hit the gratings in
the local coordinate system. This can be used to simulate
manufacturing uncertainties or intentional grating period
variation. The callable has to return a vector of lengths N
that contains the grating constant for each photon in mm.
order_selector : callable
A function or callable object that accepts arrays of photon
energy, polarization and the blaze angle as input and returns
arrays for grating order (integer) and probability
(float). The probabiliy expresses the chance that the photon
passes the grating and is not absorbed, e.g. if the
probability that a photon at energy E ends up in order=[-2,
-1, 0, 1, 2] is [0, 0, .5, .3, .0] , then the returned
probability for all photons should be 0.8.
transmission : bool
Set to ``True`` for a transmission grating and to ``False`` for a
reflection grating. (*Default*: ``True`` ) **Warning: Reflection
gratings have not yet been tested**
groove_angle : float
Angle between the local z axis and the direction of the
grooves in radian. (*Default*: ``0.``)
'''
loc_coos_name = ['grat_y', 'grat_z']
'''name for output columns that contain the interaction point in local coordinates.'''
order_name = 'order'
'''Column name for diffraction order.'''
blaze_name = 'blaze'
'''Column name for blaze angle.'''
[docs]
def order_sign_convention(self, p, e_perp_groove):
'''Set sign convention for grating orders.
This sets the following, somewhat arbitrary convention:
Positive grating orders will are displaced along the local
:math:`\hat e_z` vector, negative orders in the opposite
direction. If the grating is rotated by :math:`-\pi` the
physical situation is the same, but the sign of the grating
order will be reversed. In this sense, the convention chosen
is arbitrarily. However, it has some practical advantages: The
implementation is fast and all photons passing through the
grating in the same diffraction order are displaced in the
same way. (Contrary to the convention in :class:`CATGrating`.)
``order_sign_convention`` has to be a callable that accepts an
array of eukledian direction vectors as input and returns
``+1``, ``-1``, or an array filled with ``-1`` or ``+1``.
Parameters
----------
p : np.array
Array of Eucleadian direction vectors
e_perp_groove : np.array
Array of local groove directions at the positions where the photons
hit
'''
# -1 because n, l, d should be right-handed coordinate system
# while n = e_x, l = e_x, and d = e_y would be left-handed.
return -1
def __init__(self, **kwargs):
self.order_selector = kwargs.pop('order_selector')
self.transmission = kwargs.pop('transmission', True)
if 'd' not in kwargs:
raise ValueError('Input parameter "d" (Grating constant) is required.')
self._d = kwargs.pop('d')
groove_angle = kwargs.pop('groove_angle', 0.)
super().__init__(**kwargs)
self.geometry._geometry['groove_angle'] = groove_angle
[docs]
def e_groove_coos(self, intercoos):
'''Get coordiante orthonormal coordinate system along groove direction.
Parameters
----------
intercoos : `numpy.ndarray` of shape (N, 2)
coordinates in the coordiante system of the geometry (e.g. (x, y), or
(r, phi)).
Returns
-------
e_groove, e_perp_groove, n : `numpy.ndarray` of shape (N, 4)
Vectors pointing along the groove direction (parallel to the surface),
perpendicular to the groove direction (parallel to surface),
and normal to surface.
'''
groove = self.geometry['groove_angle']
ex, ey, en = self.geometry.get_local_euklid_bases(intercoos)
e_groove = math.sin(-groove) * ex + math.cos(-groove) * ey
e_perp_groove = math.cos(-groove) * ex - math.sin(-groove) * ey
return e_groove, e_perp_groove, en
[docs]
def d(self, intercoos):
'''Method that returns the grating constant at given positions.
For a grating with constant grating constant, this will just return the
number input as ``d`` when the element was initialized. For gratings
where the grating constant varies with position on the facet, this
calculates the appropriate number for every position.
'''
if not callable(self._d):
return self._d
else:
return self._d(intercoos)
[docs]
def blaze_angle_modifier(self, intercoos):
'''Modify blaze angle
In `diffract_photons` the blaze angle is calculated relative
to the surface of the grating. In cases where that number has
to be modified (e.g. when the grating bars are not
perpendicular to the surface) the blaze angle can be modified
by overriding this function.
'''
return np.zeros(intercoos.shape[0])
[docs]
def diffract_photons(self, photons, intersect, interpos, intercoos):
'''Vectorized implementation'''
p = norm_vector(photons['dir'].data[intersect])
l, d, n = self.e_groove_coos(intercoos[intersect])
# Minus sign here because we want n, l, d to be a right-handed coordinate system
d = - d
wave = energy2wave / photons['energy'].data[intersect]
# calculate angle between normal and (ray projected in plane perpendicular to groove)
# -> this is the blaze angle
p_perp_to_grooves = norm_vector(p - np.einsum("ij,ij->i", p, l)[:, None] * l)
# Use abs here so that blaze angle is always in 0..pi/2
# independent of the relative orientation of p and n.
# Also, clip because rounding errors can lead to products slightly larger than 1
blazeangle = np.arccos(np.clip(np.abs(np.einsum("ij,ij->i", p_perp_to_grooves,
n)), 0, 1))
blazeangle += self.blaze_angle_modifier(intercoos[intersect, :])
m, prob = self.order_selector(photons['energy'].data[intersect],
photons['polarization'].data[intersect],
blazeangle)
# The idea to calculate the components in the (d,l,n) system separately
# is taken from MARX
sign = self.order_sign_convention(p, d)
p_d = np.einsum("ij,ij->i", p, d) + sign * m * wave / self.d(intercoos[intersect, :])
p_l = np.einsum("ij,ij->i", p, l)
# The norm for p_n can be derived, but the direction needs to be chosen.
p_n = np.sqrt(1. - p_d**2 - p_l**2)
# Check if the photons have same direction compared to normal before
direction = np.sign(np.einsum("ij,ij->i", p, n), dtype=float)
if not self.transmission:
direction *= -1
dir = p_d[:, None] * d + p_l[:, None] * l + (direction * p_n)[:, None] * n
return dir, m, prob, blazeangle
[docs]
def specific_process_photons(self, photons,
intersect, interpos, intercoos):
dir, m, p, blaze = self.diffract_photons(photons, intersect, interpos,
intercoos)
pol = parallel_transport(photons['dir'].data[intersect, :], dir,
photons['polarization'].data[intersect, :])
out = {'dir': dir, 'probability': p, 'polarization': pol,
self.order_name: m, self.blaze_name: blaze}
return out
[docs]
class CATGrating(FlatGrating):
'''Critical-Angle-Transmission Grating
CAT gratings are a special case of :class:`FlatGrating` and accept
the same arguments.
They differ from a :class:`FlatGrating` in the sign convention of
the grating orders: Blazing happens on the side of the negative
orders. Obviously, this convention is only meaningful if the
photons do not arrive perpendicular to the grating.
'''
[docs]
def order_sign_convention(self, p, e_perp_groove):
'''Convention to chose the sign for CAT grating orders
Blazing happens on the side of the negative orders. Obviously, this
convention is only meaningful if the photons do not arrive
perpendicular to the grating.
'''
dotproduct = np.einsum("ij,ij->i", p, e_perp_groove)
sign = np.sign(dotproduct)
sign[sign == 0] = 1
return sign