Merge pull request #776 from yeganer/plasma_fix_series_or

Fix compatibility bug in radiative_properties
tardis-sn · Aug 17, 2017 · 8341479 · 8341479
2 parents a204bde + ef5e3ef
commit 8341479
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/tardis/plasma/properties/radiative_properties.py b/tardis/plasma/properties/radiative_properties.py
@@ -68,11 +68,11 @@ def calculate(self, g, level_number_density, lines_lower_level_index,
         stimulated_emission_factor[meta_stable_upper &
                                    (stimulated_emission_factor < 0)] = 0.0
         if self.nlte_species:
-            nlte_lines_mask = \
-                np.zeros(stimulated_emission_factor.shape[0]).astype(bool)
-            for species in self.nlte_species:
-                nlte_lines_mask |= (lines.atomic_number == species[0]) & \
-                                   (lines.ion_number == species[1])
+            nlte_lines_mask = lines.apply(
+                    lambda row:
+                    (row.atomic_number, row.ion_number) in self.nlte_species,
+                    axis=1
+            ).values
             stimulated_emission_factor[(stimulated_emission_factor < 0) &
                 nlte_lines_mask[np.newaxis].T] = 0.0
         return stimulated_emission_factor