Merge branch 'master' into 35-superposition-z

.github/workflows/lint.yml

@@ -12,7 +12,7 @@ jobs:
       - name: Setup Python
         uses: actions/setup-python@v1
         with:
-          python-version: 3.7.4
+          python-version: 3.7
           architecture: x64
       - name: Checkout PyTorch
         uses: actions/checkout@master

aisim/atoms.py

@@ -1,6 +1,7 @@
 """Classes and functions related to the atomic cloud."""
 
 import numpy as np
+import scipy.linalg as splin
 from . import convert
 
 
@@ -76,10 +77,18 @@ def __getitem__(self, key):
             a new instance of the atomic ensemble only containing the selected
             atoms
         """
+        phase_space_vectors = self.phase_space_vectors[key][:]
+        state_kets = self.state_kets[key][:]
+        weights = self.weights[key]
+        if isinstance(key, int):
+            # ratain correct shape in case of only one atom is selected
+            phase_space_vectors = phase_space_vectors.reshape(1, 6)
+            state_kets = state_kets.reshape(1, len(state_kets))
+            weights = weights.reshape(1, 1)
         new_instance = AtomicEnsemble(
-            phase_space_vectors=self.phase_space_vectors[key][:],
-            state_kets=self.state_kets[key][:],
-            weights=self.weights[key])
+            phase_space_vectors=phase_space_vectors,
+            state_kets=state_kets,
+            weights=weights)
         return new_instance
 
     def __len__(self):
@@ -198,6 +207,26 @@ def state_occupation(self, state):
         occupation = np.abs(np.matmul(projection_bras, self.state_kets))**2
         return occupation.flatten()
 
+    def fidelity(self, rho_target):
+        """
+        Calculate fidelity of ensemble's density matrix and target matrix [1].
+
+        Parameters
+        ----------
+        rho_target : array
+            target density matrix as m x m array
+
+        Returns
+        -------
+        fidelity : float
+            fidelity of AtomicEnsemble's compared to target density matrix
+
+        References
+        ----------
+        [1] https://en.wikipedia.org/wiki/Fidelity_of_quantum_states
+        """
+        return _fidelity(self.density_matrix, rho_target)
+
 
 def create_random_ensemble_from_gaussian_distribution(pos_params, vel_params,
                                                       n_samples, seed=None,
@@ -366,3 +395,34 @@ def combine_weights(pos_weights, vel_weights):
     """
     # FIXME: replace with faster version, for example based on meshgrid
     return np.array([p * v for p in pos_weights for v in vel_weights])
+
+
+def _fidelity(rho_a, rho_b):
+    """
+    Calculate the fidelity of two density matrices [1, 2].
+
+    Parameters
+    ----------
+    rho_a : array
+        density matrix as m x m array
+    rho_b : array
+        density matrix as m x m array
+
+    Returns
+    -------
+    fidelity : float
+        fidelity of both density matrices
+
+    References
+    ----------
+    [1] https://en.wikipedia.org/wiki/Fidelity_of_quantum_states
+    [2] http://qutip.org/docs/4.0.2/modules/qutip/metrics.html
+    """
+    assert rho_a.shape == rho_b.shape
+    sqrt_rho_a = splin.sqrtm(rho_a)  # matrix square root
+    # Implementation used in qutip"s fidelity calculation: calculating the
+    # eigenvalues and taking it's square root instead of matrix square root and
+    # taking the trace. It's faster.
+    eig_vals = np.linalg.eigvals(sqrt_rho_a @ rho_b @ sqrt_rho_a)
+    fidelity = np.real(np.sum(np.sqrt(eig_vals)))**2
+    return fidelity

aisim/det.py

@@ -6,28 +6,34 @@
 
 class Detector():
     """
-    A circular detection area.
+    A generic detection zone.
+
+    This is only a template without functionality. Deriving classes have to
+    implement `_detected_idx`.
 
     Parameters
     ----------
-    r_det : float
-        radius of the detection area in the x, y plane
     t_det : float
         time of the detection
+    **kwargs :
+        Additional arguments used by classes that inherit from this class. All
+        keyworded arguments are stored as attribues.
 
     Attributes
     ----------
-    r_det : float
-        radius of the detection area in the x, y plane
     t_det : float
         time of the detection
+    ** kwargs :
+        all keyworded arguments that are passed upon creation
     """
 
-    def __init__(self, r_det, t_det):
-        self.r_det = r_det
+    def __init__(self, t_det, **kwargs):
         self.t_det = t_det
+        # save all keyworded arguments as attributes
+        for key, value in kwargs.items():
+            setattr(self, key, value)
 
-    def detected_idx(self, atoms):
+    def _detected_idx(self, atoms):
         """
         Return indices of the detected atoms.
 
@@ -45,10 +51,8 @@ def detected_idx(self, atoms):
         det_idx : nd array of bool
             boolean array for filtering an AtomicEnsemble; True if detected,
             otherwise False.
-
         """
-        rho = convert.cart2pol(atoms.calc_position(self.t_det))[:, 0]
-        return np.where(rho <= self.r_det, True, False)
+        raise NotImplementedError()
 
     def detected_atoms(self, atoms):
         """
@@ -68,4 +72,59 @@ def detected_atoms(self, atoms):
             atomic ensemble containing only phase space vectors that are
             eventually detected
         """
-        return atoms[self.detected_idx(atoms)]
+        return atoms[self._detected_idx(atoms)]
+
+
+class SphericalDetector(Detector):
+    """
+    A spherical detection zone.
+
+    All atoms within a sphere of the specified radius will be detected.
+
+    Parameters
+    ----------
+    t_det : float
+        time of the detection
+    r_det :
+        radius of the spherical detection zone
+
+    Attributes
+    ----------
+    t_det : float
+        time of the detection
+    ** kwargs :
+        all keyworded arguments that are passed upon creation
+    """
+
+    def _detected_idx(self, atoms):
+        det_pos = atoms.calc_position(self.t_det)
+        x, y, z = det_pos[:, 0], det_pos[:, 1], det_pos[:, 2]
+        rho = np.sqrt(x**2 + y**2 + z**2)
+        return np.where(rho <= self.r_det, True, False)
+
+
+class PolarDetector(Detector):
+    """
+    A spherical detection zone.
+
+    All atoms within a circle with the specified radius within the x-y plane
+    will be detected.
+
+    Parameters
+    ----------
+    t_det : float
+        time of the detection
+    r_det :
+        radius of the spherical detection zone
+
+    Attributes
+    ----------
+    t_det : float
+        time of the detection
+    ** kwargs :
+        all keyworded arguments that are passed upon creation
+    """
+
+    def _detected_idx(self, atoms):
+        rho = convert.cart2pol(atoms.calc_position(self.t_det))[:, 0]
+        return np.where(rho <= self.r_det, True, False)

docs/examples/wavefront-aberrations.ipynb

@@ -53,7 +53,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -69,7 +69,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -82,7 +82,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [
     {
@@ -126,7 +126,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -162,7 +162,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -172,7 +172,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [
     {
@@ -181,7 +181,7 @@
        "Text(0, 0.5, 'y / mm')"
       ]
      },
-     "execution_count": 8,
+     "execution_count": 7,
      "metadata": {},
      "output_type": "execute_result"
     },
@@ -208,7 +208,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [
     {
@@ -217,7 +217,7 @@
        "Text(0, 0.5, 'weights')"
       ]
      },
-     "execution_count": 9,
+     "execution_count": 8,
      "metadata": {},
      "output_type": "execute_result"
     },
@@ -252,7 +252,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -270,7 +270,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 10,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -282,15 +282,15 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [],
    "source": [
     "awfs = []\n",
     "r_dets = np.linspace(2e-3, 10e-3, 10)\n",
     "for r_det in r_dets:\n",
     "    # creating detector with new detection radius\n",
-    "    det = ais.Detector(r_det, t_det)\n",
+    "    det = ais.PolarDetector(t_det, r_det=r_det)\n",
     "    \n",
     "    det_atoms = det.detected_atoms(atoms)\n",
     "    \n",
@@ -318,7 +318,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 12,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -330,16 +330,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 13,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "<matplotlib.legend.Legend at 0x261d4ec9ac8>"
+       "<matplotlib.legend.Legend at 0x1ad8144ec48>"
       ]
      },
-     "execution_count": 14,
+     "execution_count": 13,
      "metadata": {},
      "output_type": "execute_result"
     },

setup.cfg

@@ -21,6 +21,7 @@ setup_requires =
 install_requires = 
     numpy
     matplotlib
+    scipy
 packages = find: