Merge pull request #8 from ufo-kit/update-fixed-spectrum-source

Update fixed spectrum source

examples/source.py

@@ -3,8 +3,10 @@
 import numpy as np
 import quantities as q
 import syris
+import syris.imageprocessing as ip
 from syris.geometry import Trajectory
-from syris.devices.sources import BendingMagnet
+from syris.devices.sources import BendingMagnet, FixedSpectrumSource
+from .util import get_default_parser, show
 
 
 def make_triangle(n=128):
@@ -15,7 +17,7 @@ def make_triangle(n=128):
     return list(zip(z, y, z)) * q.mm
 
 
-def main():
+def run_bending_magnet():
     syris.init()
 
     ps = 1 * q.um
@@ -38,5 +40,40 @@ def main():
     plt.show()
 
 
+def run_fixed():
+    syris.init()
+    n = 512
+    ps = 1 * q.um
+    energies = np.arange(5, 30) * q.keV
+    y, x = np.mgrid[-n // 2 : n // 2, -n // 2 : n // 2]
+    flux = np.exp(-(x ** 2 + y ** 2) / (100 ** 2)) / q.s
+    weights = np.arange(1, len(energies) + 1)[:, np.newaxis, np.newaxis]
+    flux = weights * flux
+    traj = Trajectory([(n / 2, n / 2, 0)] * ps)
+    source = FixedSpectrumSource(energies, flux, 30 * q.m, (100, 500) * q.um, traj, pixel_size=ps)
+
+    im = ip.compute_intensity(source.transfer((n, n), ps, 5 * q.keV)).get()
+    show(im, title="Original sampling")
+    im = ip.compute_intensity(source.transfer((2 * n,) * 2, ps / 2, 5 * q.keV)).get()
+    show(im, title="2x supersampled")
+    plt.show()
+
+
+def main():
+    parser = get_default_parser(__doc__)
+    subparsers = parser.add_subparsers(help="sub-command help", dest="sub-commands", required=True)
+
+    bm = subparsers.add_parser("bm", help="BendingMagnet example")
+    bm.set_defaults(_func=run_bending_magnet)
+
+    fixed = subparsers.add_parser(
+        "fixed", help="FixedSpectrumSource example with Gaussian intensity profile"
+    )
+    fixed.set_defaults(_func=run_fixed)
+
+    args = parser.parse_args()
+    args._func()
+
+
 if __name__ == "__main__":
     main()

syris/bodies/simple.py

@@ -80,10 +80,11 @@ def make_sphere(n, radius, pixel_size=1 * q.m, material=None, queue=None):
     2 + 0.5, which means between two adjacent pixels. *pixel_size*, *material* and *queue*, which is
     an OpenCL command queue, are used to create :class:`.StaticBody`.
     """
+    pixel_size = make_tuple(pixel_size, num_dims=2)
     image = np.zeros((n, n), dtype=cfg.PRECISION.np_float)
     y, x = np.mgrid[-n // 2 : n // 2, -n // 2 : n // 2]
-    x = (x + 0.5) * pixel_size.simplified.magnitude
-    y = (y + 0.5) * pixel_size.simplified.magnitude
+    x = (x + 0.5) * pixel_size[1].simplified.magnitude
+    y = (y + 0.5) * pixel_size[0].simplified.magnitude
     radius = radius.simplified.magnitude
     valid = np.where(x ** 2 + y ** 2 < radius ** 2)
     image[valid] = 2 * np.sqrt(radius ** 2 - x[valid] ** 2 - y[valid] ** 2)

syris/devices/sources.py

@@ -8,12 +8,12 @@
 import numpy as np
 import quantities as q
 import quantities.constants.electron as qe
+import pyopencl as cl
 import pyopencl.array as cl_array
 from pyopencl import clmath
 from quantities.quantity import Quantity
 from quantities.constants import fine_structure_constant
 from scipy import integrate, special
-from scipy import interpolate as interp
 import syris.config as cfg
 import syris.imageprocessing as ip
 import syris.math as smath
@@ -149,16 +149,46 @@ def get_flux(self, photon_energy, vertical_angle, pixel_size):
 
 
 class FixedSpectrumSource(XRaySource):
-    def __init__(self, energies, flux, sample_distance, size, trajectory, phase_profile="plane"):
+    def __init__(
+        self,
+        energies,
+        flux,
+        sample_distance,
+        size,
+        trajectory,
+        pixel_size=None,
+        phase_profile="plane",
+    ):
+        """Source with a pre-computed *flux* at given *energies*. *flux* can be a 1D array, in which
+        case there is no spatial distribution of intensities, or it can be a 3D array, in which case
+        every 2D image *flux[i]* represents spatial intensity distribution for *energies[i]*.
+        """
         super(FixedSpectrumSource, self).__init__(
             sample_distance, size, trajectory, phase_profile=phase_profile
         )
-        self._energies = energies.rescale(q.keV).magnitude
-        self._flux = flux.rescale(1 / q.s).magnitude
-        self._tck = interp.splrep(self._energies, self._flux)
+        if len(flux) != len(energies):
+            raise XRaySourceError("Flux must have the same length as energies")
+        if flux.ndim == 3 and pixel_size is None:
+            raise XRaySourceError("pixel_size must be specified for 3D flux input")
+        self._pixel_size = make_tuple(pixel_size, num_dims=2)
+        self._energies = energies
+        self._flux = flux
 
     def get_flux(self, photon_energy, vertical_angle, pixel_size):
-        return interp.splev(photon_energy.rescale(q.keV).magnitude, self._tck) / q.s
+        """Get linearly interpolated flux at *photon_energy*."""
+        if photon_energy <= self._energies[0]:
+            flux = self._flux[0]
+        elif photon_energy >= self._energies[-1]:
+            flux = self._flux[-1]
+        else:
+            i_0 = np.where(self._energies < photon_energy)[0][-1]
+            i_1 = np.where(self._energies >= photon_energy)[0][0]
+            d_e = (self._energies[i_1] - self._energies[i_0]).simplified.magnitude
+            w_0 = (photon_energy - self._energies[i_0]).simplified.magnitude
+            w_1 = (self._energies[i_1] - photon_energy).simplified.magnitude
+            flux = (self._flux[i_0] * w_1 + self._flux[i_1] * w_0) / d_e
+
+        return flux
 
     def _transfer_real(
         self,
@@ -174,21 +204,42 @@ def _transfer_real(
         block,
         flux=1,
     ):
-        flux = self.get_flux(energy, 0 * q.rad, pixel_size).magnitude
-        super(FixedSpectrumSource, self)._transfer_real(
-            shape,
-            center,
-            pixel_size,
-            energy,
-            exponent,
-            compute_phase,
-            is_parabola,
-            out,
+        flux = (
+            self.get_flux(energy, None, pixel_size)
+            .rescale(1 / q.s)
+            .magnitude.astype(cfg.PRECISION.np_float)
+        )
+        cl_image = gutil.get_image(flux, queue=queue)
+
+        sampler = cl.Sampler(cfg.OPENCL.ctx, False, cl.addressing_mode.CLAMP, cl.filter_mode.LINEAR)
+
+        cl_center = gutil.make_vfloat3(*center)
+        cl_ps = gutil.make_vfloat2(*pixel_size.simplified.magnitude[::-1])
+        cl_input_ps = gutil.make_vfloat2(*self._pixel_size.simplified.magnitude[::-1])
+        z_sample = self.sample_distance.simplified.magnitude
+        lam = energy_to_wavelength(energy).simplified.magnitude
+        kernel = cfg.OPENCL.programs["physics"].make_flat_from_2D_profile
+
+        ev = kernel(
             queue,
-            block,
-            flux=flux,
+            shape[::-1],
+            None,
+            out.data,
+            cl_image,
+            sampler,
+            cl_center,
+            cl_ps,
+            cl_input_ps,
+            cfg.PRECISION.np_float(z_sample),
+            cfg.PRECISION.np_float(lam),
+            np.int32(exponent),
+            np.int32(compute_phase),
+            np.int32(is_parabola),
         )
 
+        if block:
+            ev.wait()
+
 
 class BendingMagnet(XRaySource):
 

syris/gpu/opencl/physics.cl

@@ -182,6 +182,35 @@ __kernel void make_flat_from_vertical_profile(__global vcomplex *output,
                      lambda, exponent, phase, parabola);
 }
 
+/*
+ * Make flat field wavefield out of a 2D profile of intensities.
+ */
+__kernel void make_flat_from_2D_profile(__global vcomplex *output,
+                                        read_only image2d_t input,
+                                        const sampler_t sampler,
+                                        const vfloat3 center,
+                                        const vfloat2 pixel_size,
+                                        const vfloat2 input_pixel_size,
+                                        const vfloat z,
+                                        const vfloat lambda,
+                                        const int exponent,
+                                        const int phase,
+                                        const int parabola) {
+    int ix = get_global_id(0);
+    int iy = get_global_id(1);
+    float input_width = (float) get_image_width (input);
+    float input_height = (float) get_image_height (input);
+    float2 factor = (float2) (pixel_size.x / input_pixel_size.x,
+                              pixel_size.y / input_pixel_size.y);
+    float x = (ix * pixel_size.x - center.x) / input_pixel_size.x + input_width / 2.0f + 0.5f * factor.x;
+    float y = (iy * pixel_size.y - center.y) / input_pixel_size.y + input_height / 2.0f + 0.5f * factor.y;
+
+    vfloat flux = read_imagef (input, sampler, (float2) (x, y)).x * factor.x * factor.y;
+
+    make_flat_field (output, ix, iy, flux, &center, &pixel_size, z, lambda,
+                     exponent, phase, parabola);
+}
+
 /*
  * Check the sampling of a transfer function
  */

syris/tests/test_gpu_imageprocessing.py

@@ -6,9 +6,9 @@
 from syris import config as cfg
 from syris import imageprocessing as ip
 from syris.math import fwnm_to_sigma
-from syris.util import get_magnitude, make_tuple
 import itertools
 from syris.tests import are_images_supported, default_syris_init, SyrisTest
+from syris.tests.util import get_gauss_2d
 
 
 def bin_cpu(image, shape):
@@ -21,29 +21,6 @@ def bin_cpu(image, shape):
     return im[:: factor[0], :: factor[1]]
 
 
-def get_gauss_2d(shape, sigma, pixel_size=None, fourier=False):
-    shape = make_tuple(shape)
-    sigma = get_magnitude(make_tuple(sigma))
-    if pixel_size is None:
-        pixel_size = (1, 1)
-    else:
-        pixel_size = get_magnitude(make_tuple(pixel_size))
-
-    if fourier:
-        i = np.fft.fftfreq(shape[1]) / pixel_size[1]
-        j = np.fft.fftfreq(shape[0]) / pixel_size[0]
-        i, j = np.meshgrid(i, j)
-
-        return np.exp(-2 * np.pi ** 2 * ((i * sigma[1]) ** 2 + (j * sigma[0]) ** 2))
-    else:
-        x = (np.arange(shape[1]) - shape[1] // 2) * pixel_size[1]
-        y = (np.arange(shape[0]) - shape[0] // 2) * pixel_size[0]
-        x, y = np.meshgrid(x, y)
-        gauss = np.exp(-(x ** 2) / (2.0 * sigma[1] ** 2) - y ** 2 / (2.0 * sigma[0] ** 2))
-
-        return np.fft.ifftshift(gauss)
-
-
 def rescale_scipy(image, factor):
     from scipy.interpolate import RectBivariateSpline
 

syris/tests/test_sources.py

@@ -1,9 +1,12 @@
 import numpy as np
 import quantities as q
-from syris.devices.sources import BendingMagnet, Wiggler, XRaySourceError
+import unittest
+from syris.devices.sources import BendingMagnet, FixedSpectrumSource, Wiggler, XRaySourceError
 from syris.geometry import Trajectory
+import syris.imageprocessing as ip
 from syris.physics import energy_to_wavelength
-from syris.tests import default_syris_init, SyrisTest
+from syris.tests import are_images_supported, default_syris_init, SyrisTest
+from syris.tests.util import get_gauss_2d
 
 
 def make_phase(n, ps, d, energy, phase_profile):
@@ -156,3 +159,67 @@ def test_wiggler(self):
         w_u = np.abs(wiggler.transfer(shape, self.ps, energy).get()) ** 2
         u = np.abs(self.source.transfer(shape, self.ps, energy).get()) ** 2
         np.testing.assert_almost_equal(w_u / u, 4)
+
+
+class TestFixedSource(SyrisTest):
+    def setUp(self):
+        # Double precision needed for spherical phase profile
+        default_syris_init()
+        self.ps = 1 * q.um
+        self.n = 128
+        self.size = (100, 100) * q.um
+        self.sample_dist = 30 * q.m
+        self.dE = 1 * q.keV
+        self.energies = np.arange(1, 39, self.dE.magnitude) * q.keV
+        self.trajectory = Trajectory([(self.n / 2, self.n / 2, 0)] * self.ps)
+
+    def test_init(self):
+        # Wrong number of flux points
+        flux = np.arange(len(self.energies) - 1) / q.s
+        with self.assertRaises(XRaySourceError):
+            FixedSpectrumSource(self.energies, flux, self.sample_dist, self.size, self.trajectory)
+
+        # No pixel size for 3D flux
+        flux = np.mgrid[: len(self.energies), : self.n, : self.n] / q.s
+        with self.assertRaises(XRaySourceError):
+            FixedSpectrumSource(self.energies, flux, self.sample_dist, self.size, self.trajectory)
+
+    def test_get_flux(self):
+        flux = np.arange(len(self.energies)) / q.s + 10 / q.s
+        source = FixedSpectrumSource(
+            self.energies, flux, self.sample_dist, self.size, self.trajectory
+        )
+        # Out of bounds values
+        self.assertAlmostEqual(source.get_flux(self.energies[0] - self.dE, None, None), flux[0])
+        self.assertAlmostEqual(source.get_flux(self.energies[-1] + self.dE, None, None), flux[-1])
+
+        # Linear interpolation
+        self.assertAlmostEqual(
+            source.get_flux(1.2 * q.keV, None, None), 0.8 * flux[0] + 0.2 * flux[1]
+        )
+
+    @unittest.skipIf(not are_images_supported(), "Images not supported")
+    def test_transfer(self):
+        def compare_sampling(gauss, factor):
+            shape = (int(factor * self.n),) * 2
+            im = ip.compute_intensity(
+                source.transfer(shape, self.ps / factor, self.energies[0], check=False)
+            ).get()
+            self.assertAlmostEqual(im.sum(), gauss.sum(), places=3)
+            hd = ip.rescale(gauss, shape).get() / factor ** 2
+            np.testing.assert_almost_equal(im, hd, decimal=5)
+
+        gauss = np.fft.fftshift(get_gauss_2d((self.n, self.n), self.n / 10.0))
+        flux = np.tile(gauss, [len(self.energies), 1, 1]) / q.s
+        source = FixedSpectrumSource(
+            self.energies, flux, self.sample_dist, self.size, self.trajectory, pixel_size=self.ps
+        )
+
+        # Same sampling
+        compare_sampling(gauss, 1)
+
+        # Upsampling
+        compare_sampling(gauss, 2)
+
+        # Downsampling
+        compare_sampling(gauss, 0.5)

syris/tests/util.py

@@ -0,0 +1,25 @@
+import numpy as np
+from syris.util import get_magnitude, make_tuple
+
+
+def get_gauss_2d(shape, sigma, pixel_size=None, fourier=False):
+    shape = make_tuple(shape)
+    sigma = get_magnitude(make_tuple(sigma))
+    if pixel_size is None:
+        pixel_size = (1, 1)
+    else:
+        pixel_size = get_magnitude(make_tuple(pixel_size))
+
+    if fourier:
+        i = np.fft.fftfreq(shape[1]) / pixel_size[1]
+        j = np.fft.fftfreq(shape[0]) / pixel_size[0]
+        i, j = np.meshgrid(i, j)
+
+        return np.exp(-2 * np.pi ** 2 * ((i * sigma[1]) ** 2 + (j * sigma[0]) ** 2))
+    else:
+        x = (np.arange(shape[1]) - shape[1] // 2) * pixel_size[1]
+        y = (np.arange(shape[0]) - shape[0] // 2) * pixel_size[0]
+        x, y = np.meshgrid(x, y)
+        gauss = np.exp(-(x ** 2) / (2.0 * sigma[1] ** 2) - y ** 2 / (2.0 * sigma[0] ** 2))
+
+        return np.fft.ifftshift(gauss)