Fixes missing complex conjugation in transition method

aisim/atoms.py

@@ -87,13 +87,33 @@ def state_vectors(self, new_state_vectors):
             assert new_state_vectors.shape[0] == len(self)
             self._state_vectors = new_state_vectors
 
+    @property
+    def density_matrices(self):
+        """
+        n × m x m array, representing the density matrix of the m level system of the n atoms.
+        These are pure states.
+        """
+        # |Psi><Psi|
+        return np.einsum('ji,il->ijl', np.conjugate(self.state_vectors).T, self.state_vectors)
+
+    @property
+    def density_matrix(self):
+        """
+        m x m array, representing the density matrix of the AtomicEnsemble's m level system
+        """
+        pure_dm = self.density_matrices
+        n_atoms = self.state_vectors.shape[0]
+        # sum over pure |Psi><Psi| and divide by N
+        return 1/n_atoms * np.einsum('ijk->jk', pure_dm)  # sum over pure state's density matrices
+
     @property
     def position(self):
         """
         n × 3 dimensional array representing the current positions (x, y, z) of the atoms in the
         ensemble
         """
         return self.phase_space_vectors[:, 0:3]
+
     @position.setter
     def position(self, new_position):
         self.phase_space_vectors[:, 0:3] = new_position

aisim/prop.py

@@ -105,7 +105,7 @@ def transition(atoms, intensity_profile, tau, wave_vectors=None, wf=None, phase_
 
     propagator = np.array([[U_ee, U_eg], [U_ge, U_gg]], dtype='complex')
     propagator = np.transpose(propagator, (2, 0, 1))
-    # U*Psi
-    atoms.state_vectors = np.einsum('ijk,ki ->ij', propagator, atoms.state_vectors.T)
+    # U*|Psi>
+    atoms.state_vectors = np.conjugate(np.einsum('ijk,ki ->ij', propagator, np.conjugate(atoms.state_vectors).T))
     atoms = free_evolution(atoms, tau)
     return atoms