feat: add support for flexible multiplicity in CP2K inputs

dpgen/generator/lib/cp2k.py

@@ -1,5 +1,127 @@
 import numpy as np
 
+atomic_numbers = {
+    "H": 1,
+    "He": 2,
+    "Li": 3,
+    "Be": 4,
+    "B": 5,
+    "C": 6,
+    "N": 7,
+    "O": 8,
+    "F": 9,
+    "Ne": 10,
+    "Na": 11,
+    "Mg": 12,
+    "Al": 13,
+    "Si": 14,
+    "P": 15,
+    "S": 16,
+    "Cl": 17,
+    "Ar": 18,
+    "K": 19,
+    "Ca": 20,
+    "Sc": 21,
+    "Ti": 22,
+    "V": 23,
+    "Cr": 24,
+    "Mn": 25,
+    "Fe": 26,
+    "Co": 27,
+    "Ni": 28,
+    "Cu": 29,
+    "Zn": 30,
+    "Ga": 31,
+    "Ge": 32,
+    "As": 33,
+    "Se": 34,
+    "Br": 35,
+    "Kr": 36,
+    "Rb": 37,
+    "Sr": 38,
+    "Y": 39,
+    "Zr": 40,
+    "Nb": 41,
+    "Mo": 42,
+    "Tc": 43,
+    "Ru": 44,
+    "Rh": 45,
+    "Pd": 46,
+    "Ag": 47,
+    "Cd": 48,
+    "In": 49,
+    "Sn": 50,
+    "Sb": 51,
+    "Te": 52,
+    "I": 53,
+    "Xe": 54,
+    "Cs": 55,
+    "Ba": 56,
+    "La": 57,
+    "Ce": 58,
+    "Pr": 59,
+    "Nd": 60,
+    "Pm": 61,
+    "Sm": 62,
+    "Eu": 63,
+    "Gd": 64,
+    "Tb": 65,
+    "Dy": 66,
+    "Ho": 67,
+    "Er": 68,
+    "Tm": 69,
+    "Yb": 70,
+    "Lu": 71,
+    "Hf": 72,
+    "Ta": 73,
+    "W": 74,
+    "Re": 75,
+    "Os": 76,
+    "Ir": 77,
+    "Pt": 78,
+    "Au": 79,
+    "Hg": 80,
+    "Tl": 81,
+    "Pb": 82,
+    "Bi": 83,
+    "Po": 84,
+    "At": 85,
+    "Rn": 86,
+    "Fr": 87,
+    "Ra": 88,
+    "Ac": 89,
+    "Th": 90,
+    "Pa": 91,
+    "U": 92,
+    "Np": 93,
+    "Pu": 94,
+    "Am": 95,
+    "Cm": 96,
+    "Bk": 97,
+    "Cf": 98,
+    "Es": 99,
+    "Fm": 100,
+    "Md": 101,
+    "No": 102,
+    "Lr": 103,
+    "Rf": 104,
+    "Db": 105,
+    "Sg": 106,
+    "Bh": 107,
+    "Hs": 108,
+    "Mt": 109,
+    "Ds": 110,
+    "Rg": 111,
+    "Cn": 112,
+    "Nh": 113,
+    "Fl": 114,
+    "Mc": 115,
+    "Lv": 116,
+    "Ts": 117,
+    "Og": 118,
+}
+
+
 default_config = {
     "GLOBAL": {"PROJECT": "DPGEN"},
     "FORCE_EVAL": {
@@ -121,6 +243,46 @@ def iterdict(d, out_list, flag=None, indent=0):
                     out_list.insert(index, " " * indent + k + " " + v)
 
 
+def calculate_multiplicity(atom_names, atom_types, charge=0):
+    """
+    Calculate the multiplicity based on atom species, quantities, and system charge.
+
+    This function provides a basic heuristic for determining multiplicity:
+    - Even number of electrons -> singlet (multiplicity = 1)
+    - Odd number of electrons -> doublet (multiplicity = 2)
+
+    Note: This approach assumes that an odd electron count always results in a doublet state.
+    It does not account for systems with multiple unpaired electrons, which can have higher
+    multiplicities (e.g., triplet, quartet, etc.). Users should be aware of this limitation
+    and use the function accordingly.
+
+    :param atom_names: List of element symbols.
+    :param atom_types: List of atom type indices.
+    :param charge: System charge (default: 0).
+    :return: Multiplicity.
+    """
+    # Calculate the total number of electrons
+    total_electrons = 0
+    for idx in atom_types:
+        element = atom_names[idx]
+        try:
+            total_electrons += atomic_numbers[element]
+        except KeyError:
+            raise ValueError(f"Unknown element '{element}' encountered in atom_names.")
+
+    # Subtract/add electrons based on system charge
+    # Positive charge means we remove electrons, negative charge means we add electrons
+    total_electrons -= charge
+
+    # Determine multiplicity based on the total number of electrons
+    # Even number of electrons -> singlet (multiplicity = 1)
+    # Odd number of electrons -> doublet (multiplicity = 2)
+    if total_electrons % 2 == 0:
+        return 1
+    else:
+        return 2
+
+
 def make_cp2k_input(sys_data, fp_params):
     # covert cell to cell string
     cell = sys_data["cells"][0]
@@ -132,14 +294,27 @@ def make_cp2k_input(sys_data, fp_params):
     cell_c = np.array2string(cell[2, :])
     cell_c = cell_c[1:-1]
 
+    atom_names = sys_data["atom_names"]
+    atom_types = sys_data["atom_types"]
+    # Get system charge if provided, default to 0
+    charge = sys_data.get("charge", 0)
+    dft_params = fp_params.get("FORCE_EVAL", {}).get("DFT", {})
+    if "MULTIPLICITY" in dft_params:
+        multiplicity = dft_params["MULTIPLICITY"]
+    else:
+        multiplicity = calculate_multiplicity(atom_names, atom_types, charge)
+
     # get update from user
     user_config = fp_params
     # get update from cell
     cell_config = {
         "FORCE_EVAL": {"SUBSYS": {"CELL": {"A": cell_a, "B": cell_b, "C": cell_c}}}
     }
+    # get update for multiplicity
+    multiplicity_config = {"FORCE_EVAL": {"DFT": {"MULTIPLICITY": multiplicity}}}
     update_dict(default_config, user_config)
     update_dict(default_config, cell_config)
+    update_dict(default_config, multiplicity_config)
     # output list
     input_str = []
     iterdict(default_config, input_str)