def calculate(self, atomic_data, continuum_interaction_species):
two_photon_data = atomic_data.two_photon_data
mask_selected_species = two_photon_data.index.droplevel(
["level_number_lower", "level_number_upper"]
).isin(continuum_interaction_species)
if not mask_selected_species.sum():
raise IncompleteAtomicData(
"two photon transition data for the requested "
"continuum_interactions species: {}".format(
continuum_interaction_species.values.tolist()
)
)
two_photon_data = two_photon_data[mask_selected_species]
index_lower = two_photon_data.index.droplevel("level_number_upper")
index_upper = two_photon_data.index.droplevel("level_number_lower")
source_idx = atomic_data.macro_atom_references.loc[
index_upper
].references_idx
destination_idx = atomic_data.macro_atom_references.loc[
index_lower
].references_idx
two_photon_idx = pd.DataFrame(
{
"source_level_idx": source_idx.values,
"destination_level_idx": destination_idx.values,
},
index=two_photon_data.index,
)
if len(two_photon_data) != 1:
raise NotImplementedError(
"Currently only one two-photon decay is supported but there "
f"are {len(two_photon_data)} in the atomic data."
)
return two_photon_data, two_photon_idx
def _filter_atomic_property(self, ionization_data, selected_atoms):
mask = ionization_data.index.isin(selected_atoms, level="atomic_number")
ionization_data = ionization_data[mask]
counts = ionization_data.groupby(level="atomic_number").count()
if np.alltrue(counts.index == counts):
return ionization_data
else:
raise IncompleteAtomicData(
f"ionization data for the ion ({str(counts.index[counts.index != counts])}, {str(counts[counts.index != counts])})"
)
def _filter_atomic_property(self, ionization_data, selected_atoms):
ionization_data['atomic_number'] = ionization_data.index.labels[0] + 1
ionization_data['ion_number'] = ionization_data.index.labels[1] + 1
ionization_data = ionization_data[ionization_data.atomic_number.isin(
selected_atoms)]
ion_data_check = counter(ionization_data.atomic_number.values)
keys = np.array(ion_data_check.keys())
values = np.array(ion_data_check.values())
if np.alltrue(keys == values):
return ionization_data
else:
raise IncompleteAtomicData('ionization data for the ion (' +
str(keys[keys != values]) +
str(values[keys != values]) + ')')