1 files changed, 90 insertions, 0 deletions
diff --git a/Source/Utils/Physics/write_atomic_data_cpp.py b/Source/Utils/Physics/write_atomic_data_cpp.py
new file mode 100644
index 000000000..11cd3b2c0
--- /dev/null
+++ b/Source/Utils/Physics/write_atomic_data_cpp.py
@@ -0,0 +1,90 @@
+#!/usr/bin/env python3
+#
+# Copyright 2019-2021 Axel Huebl, Luca Fedeli, Maxence Thevenet
+#
+#
+# This file is part of WarpX.
+#
+# License: BSD-3-Clause-LBNL
+
+'''
+This python script reads ionization tables in atomic_data.txt (generated from
+the NIST website) and extracts ionization levels into C++ file
+IonizationEnergiesTable.H, which contains tables + metadata.
+'''
+
+import os
+import re
+
+import numpy as np
+
+filename = os.path.join( '.', 'atomic_data.txt' )
+with open(filename) as f:
+    text_data = f.read()
+
+# Read full table from file and get names, atomic numbers and offsets
+# position in table of ionization energies for all species
+regex_command = '\n\s+(\d+)\s+\|\s+([A-Z]+[a-z]*)\s+\w+\s+\|\s+\+*(\d+)\s+\|\s+\(*\[*(\d+\.*\d*)'
+list_of_tuples = re.findall( regex_command, text_data )
+ion_atom_numbers = [int(i) for i in list(dict.fromkeys( [x[0] for x in list_of_tuples] ))]
+ion_names = list(dict.fromkeys( [x[1] for x in list_of_tuples] ))
+ion_offsets = np.concatenate(([0], np.cumsum(np.array(ion_atom_numbers)[:-1])), axis=0)
+
+# Head of CPP file
+cpp_string = '// This script was automatically generated!\n'
+cpp_string += '// Edit dev/Source/Utils/Physics/write_atomic_data_cpp.py instead!\n\n'
+cpp_string += '#ifndef WARPX_UTILS_PHYSICS_IONIZATION_TABLE_H_\n'
+cpp_string += '#define WARPX_UTILS_PHYSICS_IONIZATION_TABLE_H_\n\n'
+cpp_string += '#include <AMReX_REAL.H>\n\n'
+cpp_string += '#include <map>\n'
+cpp_string += '#include <string>\n\n'
+cpp_string += 'namespace utils::physics\n'
+cpp_string += '{\n'
+
+# Map each element to ID in table
+cpp_string += '    static std::map<std::string, int> const ion_map_ids = {'
+for count, name in enumerate(ion_names):
+    cpp_string += '\n        {"' + name + '", ' + str(count) + '},'
+cpp_string = cpp_string[:-1]
+cpp_string += ' };\n\n'
+
+# Atomic number of each species
+cpp_string += '    constexpr int nelements = ' + str(len(ion_names)) + ';\n\n'
+cpp_string += '    constexpr int ion_atomic_numbers[nelements] = {\n        '
+for count, atom_num in enumerate(ion_atom_numbers):
+    if count%10==0 and count>0: cpp_string = cpp_string[:-2] + ',\n        '
+    cpp_string += str(atom_num) + ', '
+cpp_string = cpp_string[:-2]
+cpp_string += '};\n\n'
+
+# Offset of each element in table of ionization energies
+cpp_string += '    constexpr int ion_energy_offsets[nelements] = {\n        '
+for count, offset in enumerate(ion_offsets):
+    if count%10==0 and count>0: cpp_string = cpp_string[:-2] + ',\n        '
+    cpp_string += str(offset) + ', '
+cpp_string = cpp_string[:-2]
+cpp_string += '};\n\n'
+
+# Table of ionization energies
+cpp_string += '    constexpr int energies_tab_length = ' + str(len(list_of_tuples)) + ';\n\n'
+cpp_string += '    constexpr amrex::Real table_ionization_energies[energies_tab_length]{'
+for element in ion_names:
+    cpp_string += '\n        // ' + element + '\n        '
+    regex_command = \
+        '\n\s+(\d+)\s+\|\s+%s\s+\w+\s+\|\s+\+*(\d+)\s+\|\s+\(*\[*(\d+\.*\d*)' \
+        %element
+    list_of_tuples = re.findall( regex_command, text_data )
+    for count, energy in enumerate([x[2] for x in list_of_tuples]):
+        if count%3==0 and count>0: cpp_string = cpp_string[:-2] + ',\n        '
+        cpp_string += "amrex::Real(" + energy + '), '
+    cpp_string = cpp_string[:-1]
+cpp_string = cpp_string[:-1]
+cpp_string += '\n    };\n\n'
+
+cpp_string += '}\n\n'
+
+# Write the string to file
+cpp_string += '#endif // #ifndef WARPX_UTILS_PHYSICS_IONIZATION_TABLE_H_\n'
+f= open("IonizationEnergiesTable.H","w")
+f.write(cpp_string)
+f.close()