Source code for utils

import json
import numpy as np
import tensorflow as tf
import diabfunct
from typing import Union, List, Optional, Tuple
from scipy.spatial.transform import Rotation as R
from itertools import combinations_with_replacement
from dataclasses import dataclass, field

# Atomic weights of elements
atomic_weights = {
    "H": 1.00782503223,            # Hydrogen
    "He": 4.00260325413,           # Helium
    "Li": 6.938,                   # Lithium (isotopic composition varies, so this is an approximate value)
    "Be": 9.0121831,               # Beryllium
    "B": 10.806,                   # Boron (range: 10.806–10.821)
    "C": 12.0096,                  # Carbon (range: 12.0096–12.0116)
    "N": 14.00643,                 # Nitrogen (range: 14.00643–14.00728)
    "O": 15.99903,                 # Oxygen (range: 15.99903–15.99977)
    "F": 18.998403163,             # Fluorine
    "Ne": 20.1797,                 # Neon (range: 20.1797–20.1798)
    "Na": 22.98976928,             # Sodium
    "Mg": 24.304,                  # Magnesium (range: 24.304–24.307)
    "Al": 26.9815385,              # Aluminum
    "Si": 28.084,                  # Silicon (range: 28.084–28.086)
    "P": 30.973761998,             # Phosphorus
    "S": 32.059,                   # Sulfur (range: 32.059–32.076)
    "Cl": 35.446,                  # Chlorine (range: 35.446–35.457)
    "Ar": 39.792                   # Argon (range: 39.792–39.963)
}

[docs]def read_xyz(filename, set_COM=True): """ Reads an .xyz file and returns the atomic data. Args: filename (str): Path to the .xyz file. set_COM (bool): If True, centers the coordinates at the center of mass. Returns: list: A list of dictionaries, where each dictionary contains the atomic symbol and its coordinates. """ atoms = [] with open(filename, 'r') as file: # Read the number of atoms (first line) num_atoms = int(file.readline().strip()) # Read the comment/title line (second line) comment = file.readline().strip() # Read the cartesian coordinates for _ in range(num_atoms): line = file.readline().strip().split() if len(line) == 4: atom, x, y, z = line atoms.append({ 'atom': atom.capitalize(), 'xyz': [float(x), float(y), float(z)] }) else: raise ValueError("Invalid .xyz file format: expected 4 columns per atom line.") if set_COM: # Compute the center of mass coords = np.array([atom["xyz"] for atom in atoms]) # Use the atomic weights to compute the center of mass masses = np.array([atomic_weights[atom["atom"].capitalize()] for atom in atoms]) total_mass = np.sum(masses) COM = np.sum(coords * masses[:, np.newaxis], axis=0) / total_mass # Center the coordinates at the COM for atom in atoms: atom["xyz"] = list(np.array(atom["xyz"]) - COM) # Print warning if the initial geometry was not centered at the COM if not np.allclose(COM, 0.0): print(f"Warning: The geometry was not centered at the center of mass. COM: {COM}") else: print("The geometry was already centered at the center of mass.") return atoms
# Custom JSON Encoder to handle various data types
[docs]class CustomEncoder(json.JSONEncoder):
[docs] def default(self, obj): if isinstance(obj, (tf.Tensor, np.ndarray)): return obj.tolist() elif isinstance(obj, (np.floating, float)): return float(obj) elif isinstance(obj, (np.integer, int)): return int(obj) elif isinstance(obj, (np.bool_, bool)): return bool(obj) elif isinstance(obj, complex): return [obj.real, obj.imag] else: return super().default(obj)
# Helper function to process parameters
[docs]def process_parameters(param_name, vcham_params, count, n_var): """ Extract parameters from a tf.Variable object in vcham_params by name. :param param_name: Name of the parameter to find (e.g., 'funct_param:0'). :param vcham_params: List of tf.Variable objects for the current mode. :param count: Current index for slicing the parameter data. :param n_var: Number of variables to extract. :return: Sliced parameter data and updated count. """ for param in vcham_params: if param.name.startswith(param_name): return param.numpy()[count:count + n_var], count + n_var return None, count
# raise ValueError(f"Parameter '{param_name}' not found in vcham_params.") # Convert VCSystem to JSON
[docs]def VCSystem_to_json(VCSystem, general_data: dict = {}, output_name: str = "vcham_data.json", rewrite = False): verbose = False data = [] lambdas_indexes = VCSystem.idx_dict["lambda"] jt_indexes_off = VCSystem.idx_dict["jt_off"] summary_output = VCSystem.summary_output vcham_param = VCSystem.optimized_params for mode in range(VCSystem.number_normal_modes): mode_data = { "mode": mode, "diagonal": [], "non-diagonal": None } if verbose: print(f"--------------- Mode {mode} ----------------") counts = {"fn": 0, "kappa": 0, "lambda": 0} # Track counts for different parameters for state in range(VCSystem.number_states): function_type = VCSystem.diab_funct[mode][state] n_var = diabfunct.n_var[function_type] if verbose: print(f"--------------- State {state} ----------------") print(f"Nvar: {n_var}") state_data = { "state": state, "diab_funct": function_type, } inactive_mode = False if summary_output[mode][state] == "": parameters, counts["fn"] = process_parameters( "funct_param:0", vcham_param[mode], counts["fn"], n_var ) state_data["parameters"] = parameters # if mode == 8: # print(f"Parameters Function: {parameters}") # if verbose: # print(f"Parameters Function: {parameters}") elif summary_output[mode][state] == "JT": parameters, counts["fn"] = process_parameters( "jt_on_param:0", vcham_param[mode], counts["fn"], n_var ) # print(parameters) # print(f"JT parameters in mode {mode} state {state}: {parameters}") if parameters is None or len(parameters) == 0: # print(f"Empty Function parameters in mode {mode} state {state}") parameters, counts["fn"] = process_parameters( "jt_on_param_inactive:0", vcham_param[mode], counts["fn"], n_var ) inactive_mode = True if parameters is not None and len(parameters) > 0 else False # if again parameters are empty, then take the info from the previous state if parameters is None or len(parameters) == 0: # print(f"Empty JT parameters in mode {mode} state {state}") # print(F"mode_data: {mode_data}") parameters = mode_data["diagonal"][state-1]["parameters"] state_data["parameters"] = parameters state_data["parameters"] = parameters # print(f"Inactive mode: {inactive_mode}") if inactive_mode: parameters, counts["fn"] = process_parameters( "jt_kappa_param_inactive:0", vcham_param[mode], counts["kappa"], n_var ) # print(f"JT kappa parameters in mode {mode} state {state}: {parameters}") state_data["kappa"] = parameters[0] # if kappa from the previous state is non-zero, then take it if state > 0: if mode_data["diagonal"][state-1].get("kappa") is not None: state_data["kappa"] = - mode_data["diagonal"][state-1].get("kappa") elif summary_output[mode][state] == "kappa": # print("Kappa") kappa, counts["kappa"] = process_parameters( "kappa_param:0", vcham_param[mode], counts["kappa"], 1 ) parameters, counts["fn"] = process_parameters( "funct_param:0", vcham_param[mode], counts["fn"], n_var ) state_data["parameters"] = parameters state_data["kappa"] = kappa[0] if verbose: print(f"Parameters Function: {parameters}") print(f"Parameters Kappa: {kappa}") mode_data["diagonal"].append(state_data) # This only goes one time per mode for param in vcham_param[mode]: if param.name == "lambda_param:0": lambd = param.numpy() non_diagonal = { "idx": str(lambdas_indexes[mode]), "lambda": lambd } if verbose: print(f"idx: {lambdas_indexes[mode]}") print(f"Parameters Lambda: {lambd}") print(f"Non-diagonal: {non_diagonal}") mode_data["non-diagonal"] = non_diagonal if param.name == "jt_off_param:0": lambd = param.numpy() non_diagonal = { "idx": str(jt_indexes_off[mode]), "lambda": lambd } if verbose: print(f"idx: {jt_indexes_off[mode]}") print(f"Parameters JT-off: {lambd}") print(f"Non-diagonal: {non_diagonal}") mode_data["non-diagonal"] = non_diagonal data.append(mode_data) VCSystem.lvc_data["vcham"] = data output_data = { "general_data": general_data, "vcham_data": VCSystem.lvc_data, } # Save to JSON # Check if the output file already exists and change the name if necessary for not overwriting # if rewrite: if Path(output_name).is_file(): if rewrite: print(f"Warning: Overwriting existing file {output_name}") else: output_name = output_name.replace(".json", "_new.json") with open(output_name, "w") as json_file: json.dump(output_data, json_file, indent=2, cls=CustomEncoder) print(f"Data successfully saved to {output_name}")
import logging from pathlib import Path logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__)
[docs]@dataclass class MoleculeDipoleSettings: mol_idx: int selected_states: Optional[List[int]] = None CM: np.ndarray = field(default_factory=lambda: np.array([0.0, 0.0, 0.0])) rot_angles: Tuple[float, float, float] = (0.0, 0.0, 0.0)
[docs]@dataclass class DipoleConfig: molecules: List["MoleculeDipoleSettings"] = field(default_factory=list) selected_pairs: Optional[List[Tuple[int, int]]] = None """ If specified, only these pairs (i,j) will be computed. (i and j are indices in `molecules`.) Otherwise, compute all pairs. """
def _read_data_blocks(infiles: Union[str, List[str]]) -> List[dict]: """ Reads each JSON file from 'infiles' and returns a list of 'vcham_data' blocks. Also checks units for potential mismatches (logs a warning if found). """ if isinstance(infiles, str): infiles = [infiles] data_blocks = [] all_units = set() for path in infiles: with open(path, "r") as f: json_data = json.load(f) vcham = json_data["vcham_data"] data_blocks.append(vcham) all_units.add(vcham["units"]) # Warn if multiple units are found (optional). if len(all_units) > 1: logger.warning( f"Detected multiple units: {all_units}. " "You may need to convert them consistently before writing." ) return data_blocks def _compute_totals(data_blocks: List[dict]) -> (int, int, str): """ Compute the total number of modes and states across all data blocks. Returns (total_modes, total_states, master_units). """ total_modes = sum(block["number_normal_modes"] for block in data_blocks) total_states = sum(block["number_states"] for block in data_blocks) # We'll pick the units from the first block arbitrarily: master_units = data_blocks[0]["units"] logger.info(f"Total modes: {total_modes}, total states: {total_states}") return total_modes, total_states, master_units def _write_header_section(fh) -> None: """ Writes the MCTDH header section to the file handle 'fh'. """ fh.write( "OP_DEFINE-SECTION\n" "\tTITLE\n" "\t\tMCTDH-Operator-file created by PyVCHAM\n" "\tEND-TITLE\n" "END-OP_DEFINE-SECTION\n\n" ) def _write_parameter_section( fh, data_blocks: List[dict], master_units: str, dipoles = False, dipole_config: DipoleConfig = None ) -> None: """ Writes the PARAMETER-SECTION: - Frequencies (omega_i) - Energies (E_i) - Diabatic curves with parameters - kappa/lambda - Placeholders (gamma, iota, mu) - Dipole interaction section (optional) """ fh.write("PARAMETER-SECTION\n\n# frequencies\n") # First pass: write frequencies for each block mode_offset = 0 for block in data_blocks: nmodes = block["number_normal_modes"] vib_freqs = block["vib_freq"] # Write out each frequency with an offset for i in range(nmodes): freq = vib_freqs[i] fh.write(f"omega_{mode_offset + i + 1}\t=\t{freq:.8f} , {master_units}\n") mode_offset += nmodes fh.write("\n# energies\n") # Second pass: write energies for each block for idx_mol, block in enumerate(data_blocks): nstates = block["number_states"] e_shifts = block["energy_shift"] for s in range(nstates): E = e_shifts[s] fh.write(f"E_{idx_mol+1}_{s + 1}\t=\t{E:.8f} , {master_units}\n") fh.write("\n# Diabatic curves with parameters\n") # Third pass: write diabetic curves + kappa + lambda current_mode_offset = 0 current_state_offset = 0 for block in data_blocks: nmodes = block["number_normal_modes"] nstates = block["number_states"] vcham_modes = block["vcham"] # Diabatic curves for m in range(nmodes): diag_list = vcham_modes[m]["diagonal"] for s in range(nstates): diag = diag_list[s] diab_funct = diag["diab_funct"] parameters = diag["parameters"] for idx_param, param in enumerate(parameters): global_mode_index = current_mode_offset + m line = ( f"M{global_mode_index}S{s}_{idx_param+1}" f"\t=\t{param:.8f}" ) # If first param or function is ho/quartic => add units if idx_param == 0 or diab_funct in ["ho", "quartic"]: fh.write(line + f" , {master_units}\n") else: fh.write(line + "\n") fh.write("\n") # on-diagonal kappa fh.write("# on-diagonal linear coupling constants (kappa)\n") for m in range(nmodes): diag_list = vcham_modes[m]["diagonal"] for s in range(nstates): kappa = diag_list[s].get("kappa") if kappa is not None: global_mode_index = current_mode_offset + m # global_state_index = current_state_offset + s fh.write( f"kappa{global_mode_index + 1}_{s + 1}" f"\t=\t{kappa:.8f} , {master_units}\n" ) fh.write("\n# off-diagonal linear coupling constants (lambda)\n") for m in range(nmodes): non_diag = vcham_modes[m]["non-diagonal"] if non_diag: indexes = json.loads(non_diag["idx"]) lambdas = non_diag["lambda"] for idx_pair, (st1, st2) in enumerate(indexes): global_mode_index = current_mode_offset + m # global_st1 = current_state_offset + st1 # global_st2 = current_state_offset + st2 lam_val = lambdas[idx_pair] fh.write( f"lambda{global_mode_index + 1}_{st1 + 1}_{st2 + 1}" f"\t=\t{lam_val:.8f} , {master_units}\n" ) fh.write("\n") current_mode_offset += nmodes current_state_offset += nstates # Dipole interaction section if dipoles: # Check reference geometry ref_geometry = data_blocks[0]["reference_geometry"] if ref_geometry is None: raise ValueError("dipole=True but no reference geometry was provided!") if dipole_config is None: raise ValueError("dipole=True but no dipole_config was provided!") # Instead of calling a two-molecule function directly, # we call a new function that handles multiple molecules: dip_values = _calculate_dipoles_multi(fh, data_blocks, dipole_config) print("Dipole interactions written to file.") fh.write("\n") return dip_values # placeholders fh.write("\n# on-diagonal bilinear coupling constants (gamma)\n") fh.write("# Third order quad-lin and cube terms (iota)\n") fh.write("# off-diagonal bilinear coupling constants (mu)\n") fh.write("end-parameter-section\n\n") # Dipole interaction section return None # ----------------------------------------------------------- from itertools import combinations def _calculate_dipoles_multi(fh, data_blocks, dipole_config: DipoleConfig) -> None: """ Compute dipole-dipole interactions for either: 1. All distinct molecule pairs, if dipole_config.selected_pairs is None 2. Or only the specified pairs in dipole_config.selected_pairs """ all_molecules = dipole_config.molecules num_molecules = len(all_molecules) print(f"Computing dipole interactions for {num_molecules} molecules.") print(f"Reference geometry: {all_molecules[0].CM}") if num_molecules < 2: fh.write("# Less than 2 molecules; no pairwise dipole interactions.\n") return # If the user provided a list of pairs, use that if dipole_config.selected_pairs is not None: pair_list = dipole_config.selected_pairs else: # Otherwise, default to all unique pairs pair_list = list(combinations(range(num_molecules), 2)) fh.write("# Computing dipole interactions for selected molecule pairs.\n") fh.write(f"# Pairs: {pair_list}\n") dip_val_pairs = [] for (i, j) in pair_list: if i < 0 or i >= num_molecules or j < 0 or j >= num_molecules: raise ValueError(f"Invalid pair indices: ({i}, {j}). Out of range.") settings_i = all_molecules[i] settings_j = all_molecules[j] dip_val_pairs.append(_calculate_dipoles_pair(fh, data_blocks, settings_i, settings_j)) dip_val_pairs = np.array(dip_val_pairs) return dip_val_pairs def _calculate_dipoles_pair( fh, data_blocks: List[dict], settings_i: MoleculeDipoleSettings, settings_j: MoleculeDipoleSettings ) -> None: """ Compute dipole interactions BETWEEN TWO molecules (settings_i, settings_j). This is the same logic you had, but references: - settings_i.mol_idx - settings_i.selected_states - settings_i.CM - settings_i.rot_angles etc. """ mol1_idx = settings_i.mol_idx mol2_idx = settings_j.mol_idx # Validate that these indices exist if mol1_idx >= len(data_blocks) or mol2_idx >= len(data_blocks): raise ValueError(f"Invalid molecule indices {mol1_idx} or {mol2_idx} for dipole calc.") if "dipole_matrix" not in data_blocks[mol1_idx] or "dipole_matrix" not in data_blocks[mol2_idx]: raise ValueError("Dipole matrix not found in the relevant data blocks.") # Extract data for each molecule nstates1 = data_blocks[mol1_idx]["number_states"] nstates2 = data_blocks[mol2_idx]["number_states"] dipoles1 = data_blocks[mol1_idx]["dipole_matrix"] dipoles2 = data_blocks[mol2_idx]["dipole_matrix"] # If user doesn't specify subsets, default to all states sel_states1 = settings_i.selected_states or range(nstates1) sel_states2 = settings_j.selected_states or range(nstates2) # Precompute rotated dipoles for molecule 1 from itertools import combinations_with_replacement rotated_dipoles1 = [] indices1 = [] for s1, s2 in combinations_with_replacement(sel_states1, 2): dipole1 = dipoles1[s1][s2] mu_mol1_rot = rotate_dipole(dipole1, *settings_i.rot_angles) rotated_dipoles1.append(mu_mol1_rot) indices1.append((s1, s2)) rotated_dipoles1 = np.array(rotated_dipoles1) # Precompute rotated dipoles for molecule 2 rotated_dipoles2 = [] indices2 = [] for s3, s4 in combinations_with_replacement(sel_states2, 2): dipole2 = dipoles2[s3][s4] mu_mol2_rot = rotate_dipole(dipole2, *settings_j.rot_angles) rotated_dipoles2.append(mu_mol2_rot) indices2.append((s3, s4)) rotated_dipoles2 = np.array(rotated_dipoles2) # Compute all interactions results = [] values = [] CM1 = settings_i.CM CM2 = settings_j.CM for i1, mu1 in enumerate(rotated_dipoles1): s1, s2 = indices1[i1] for i2, mu2 in enumerate(rotated_dipoles2): s3, s4 = indices2[i2] val = dipole_interaction(mu1, mu2, CM1, CM2) values.append(val) results.append(f"DD_{mol1_idx}_{s1+1}_{s2+1}_VS_{mol2_idx}_{s3+1}_{s4+1} = {val:.8f}") # Write results fh.write(f"# Dipole interactions: Mol {mol1_idx} vs Mol {mol2_idx}\n") fh.write("\n".join(results)) fh.write("\n\n") return values
[docs]def rotate_dipole(dipole, alpha, beta, gamma): """ Rotate a dipole moment using given Euler angles. Parameters: dipole : np.array Original dipole moment vector (3D) alpha, beta, gamma : float Euler angles in degrees Returns: np.array Rotated dipole moment """ # Convert degrees to radians euler_angles = np.radians([alpha, beta, gamma]) # Create rotation object using ZYX convention (intrinsic rotation) rotation = R.from_euler('ZYX', euler_angles) # Apply rotation to dipole moment rotated_dipole = rotation.apply(dipole) return rotated_dipole
[docs]def dipole_interaction(mu_mol1, mu_mol2, CM_mol1, CM_mol2): """ Calculate dipole-dipole interaction term based on the formula from the paper J. Phys. Chem. Lett. 2023,14,11367-11375 https://doi.org/10.1021/acs.jpclett.3c02935 Parameters: mu_mol1 : np.array Dipole moment vector of molecule 1 mu_mol2 : np.array Dipole moment vector of molecule 2 CM_mol1 : np.array Center of mass coordinates of molecule 1 CM_mol2 : np.array Center of mass coordinates of molecule 2 Returns: float Interaction energy """ # H_{mol1mol2} = prefactor * (mu_mol1 * mu_mol2 - 3*(mu_mol1 * r) * (mu_mol2 * r)) # Normalize the displacement vector to get unit vector u_R CM_mol1 = np.array(CM_mol1) CM_mol2 = np.array(CM_mol2) r_vector = CM_mol2 - CM_mol1 R = np.linalg.norm(r_vector) u_R = r_vector / R # Calculate dot products dot_mu_mol1_mu_mol2 = np.dot(mu_mol1, mu_mol2) dot_mu_mol1_uR = np.dot(mu_mol1, u_R) dot_mu_mol2_uR = np.dot(mu_mol2, u_R) # Compute the interaction energy using the formula interaction = (1 / R**3) * (dot_mu_mol1_mu_mol2 - 3 * dot_mu_mol1_uR * dot_mu_mol2_uR) return interaction
def _write_labels_section(fh, data_blocks: List[dict]) -> None: """ Writes the LABELS-SECTION for all blocks (Morse/Anti-Morse). """ fh.write("LABELS-SECTION\n# Diabatic function labels\n") current_mode_offset = 0 current_state_offset = 0 for idx_mol, block in enumerate(data_blocks): nmodes = block["number_normal_modes"] nstates = block["number_states"] vcham_modes = block["vcham"] for m in range(nmodes): diag_list = vcham_modes[m]["diagonal"] for s in range(nstates): diab_funct = diag_list[s]["diab_funct"] if diab_funct in ["morse", "antimorse"]: global_mode_index = current_mode_offset + m # global_state_index = current_state_offset + s fh.write( f"v_{global_mode_index + 1}_{s + 1}" "\t=\tmorse1[" f"M{global_mode_index}S{s}_1," f"M{global_mode_index}S{s}_2," f"M{global_mode_index}S{s}_3," f"E_{idx_mol+1}_{s + 1}" "]\n" ) current_mode_offset += nmodes current_state_offset += nstates fh.write("end-labels-section\n\n") def _write_hamiltonian_section( fh, data_blocks: List[dict], total_modes: int, total_states: int, dipole=False, dipole_config: DipoleConfig = None, values_dipoles = None ) -> None: """ Writes the HAMILTONIAN-SECTION for all blocks, iterating in the same order with offset logic. """ fh.write("HAMILTONIAN-SECTION\n") fh.write("------------------------------------------------------------------\n") # e.g. "modes | v1 | v2 | v3 ... v(total_modes) | el" max_length = 60 # Maximum line length mode_header = " ".join(f"v{m+1} |" for m in range(total_modes)) header_prefix = "modes | " # Start with the prefix current_line = header_prefix lines = [] # Split based on " | " to ensure proper splitting without adding "|" segments = mode_header.split(" | ") for i, segment in enumerate(segments): if len(current_line) + len(segment) > max_length: # Save the current line and start a new one with "modes | " prefix lines.append(current_line.strip()) current_line = "modes | " + segment else: if i > 0: # Add separator only between segments, not at the start current_line += " | " current_line += segment # Add the final line with "el" appended if len(data_blocks) == 1: lines.append(current_line.strip() + " el") else: lines.append(current_line.strip() + " el1") current_line = header_prefix elec = "" if len(data_blocks) > 1: for mol in range(len(data_blocks)): if mol == len(data_blocks) - 1: elec += f" el{mol+1}" elif mol == 0: pass else: elec += f" el{mol+1} |" lines.append(current_line.strip() + elec) # Write each formatted line to the file for line in lines: fh.write(line + "\n") fh.write("------------------------------------------------------------------\n\n") # Kinetic Energy fh.write("# Kinetic Energy\n") for mode_idx in range(total_modes): fh.write(f"omega_{mode_idx + 1}\t|{mode_idx + 1}\tKE\n") # Potential & couplings current_mode_offset = 0 current_state_offset = 0 for cur_mol, block in enumerate(data_blocks): idx_mol = cur_mol + 1 nmodes = block["number_normal_modes"] nstates = block["number_states"] vcham_modes = block["vcham"] fh.write(f"\n# Molecule {cur_mol} \n") # 1. Harmonic/quartic potential fh.write("\n# Potential for Harmonic oscillator\n") for m in range(nmodes): diag_list = vcham_modes[m]["diagonal"] diab_functs = [d["diab_funct"] for d in diag_list] if any(f in ["ho", "quartic"] for f in diab_functs): global_mode_idx = current_mode_offset + m fh.write(f"0.5*omega_{global_mode_idx + 1}\t|{global_mode_idx + 1}\tq^2\n") # 2. Electronic states fh.write("\n# Electronic States\n") for s in range(nstates): fh.write( f"E_{idx_mol}_{s + 1}\t|{total_modes + 1 + cur_mol}\tS{s + 1}&{s + 1}\n" ) # 3. Lambda fh.write("\n# Lambda\n") for m in range(nmodes): non_diag = vcham_modes[m]["non-diagonal"] if non_diag: indexes = json.loads(non_diag["idx"]) lambdas = non_diag["lambda"] for i, (st1, st2) in enumerate(indexes): global_mode_idx = current_mode_offset + m fh.write( f"lambda{global_mode_idx + 1}_{st1 + 1}_{st2 + 1}" f"\t|{global_mode_idx + 1}\tq\t|{total_modes + 1 + cur_mol}" f"\tS{st1 + 1}&{st2 + 1}\n" ) # 4. Kappa fh.write("\n# Kappa\n") for m in range(nmodes): diag_list = vcham_modes[m]["diagonal"] for s in range(nstates): kappa = diag_list[s].get("kappa") if kappa is not None: global_mode_idx = current_mode_offset + m fh.write( f"kappa{global_mode_idx + 1}_{s + 1}" f"\t|{global_mode_idx + 1}\tq\t|{total_modes + 1 + cur_mol}" f"\tS{s + 1}&{s + 1}\n" ) # 5. Quartic fh.write("\n# Quartic potential\n") for m in range(nmodes): diag_list = vcham_modes[m]["diagonal"] for s in range(nstates): diag = diag_list[s] if diag["diab_funct"] == "quartic": global_mode_idx = current_mode_offset + m fh.write( f"{0.5:.8f}*M{global_mode_idx}S{s}_1" f"\t|{global_mode_idx + 1}\tq^2" f"\t|{total_modes + 1 + cur_mol}\tS{s + 1}&{s + 1}\n" ) fh.write( f"{(1/24):.8f}*M{global_mode_idx}S{s}_2" f"\t|{global_mode_idx + 1}\tq^4" f"\t|{total_modes + 1 + cur_mol}\tS{s + 1}&{s + 1}\n" ) # 6. Morse/Anti-Morse fh.write("\n# Morse/Anti-Morse potential\n") for m in range(nmodes): diag_list = vcham_modes[m]["diagonal"] for s in range(nstates): diag = diag_list[s] if diag["diab_funct"] in ["morse", "antimorse"]: global_mode_idx = current_mode_offset + m # global_state_idx = current_state_offset + s fh.write( f"1.0\t|{global_mode_idx + 1}\t" f"v_{global_mode_idx + 1}_{s + 1}" f"\t|{total_modes + 1 + cur_mol}\tS{s + 1}&{s + 1}\n" ) # 7. Dipole interaction section if dipole: fh.write("\n# Dipole interaction\n") # If the user provided a list of pairs, use that all_molecules = dipole_config.molecules num_molecules = len(all_molecules) if dipole_config.selected_pairs is not None: pair_list = dipole_config.selected_pairs else: # Otherwise, default to all unique pairs pair_list = list(combinations(range(num_molecules), 2)) for idx, (mol1, mol2) in enumerate(pair_list): # If user doesn't specify subsets, default to all states sel_states1 = all_molecules[mol1].selected_states or range(data_blocks[mol1]["number_states"]) sel_states2 = all_molecules[mol2].selected_states or range(data_blocks[mol2]["number_states"]) count = 0 for s1, s2 in combinations_with_replacement(sel_states1, 2): for s3, s4 in combinations_with_replacement(sel_states2, 2): fh.write( f"DD_{mol1}_{s1+1}_{s2+1}_VS_{mol2}_{s3+1}_{s4+1} | " f"{total_modes + 1 + mol1} S{s1+1}&{s2+1} | " f"{total_modes + 1 + mol2} S{s3+1}&{s4+1}\n") count += 1 current_mode_offset += nmodes current_state_offset += nstates fh.write("\nEND-HAMILTONIAN-SECTION\n\n") fh.write("END-OPERATOR\n")
[docs]def jsons_to_mctdh(infiles: Union[str, List[str]], outfile: str, dipole: bool = False, dipole_config: DipoleConfig = None) -> None: """ Reads multiple JSON files and writes a single MCTDH operator file that concatenates them directly (with proper indexing offsets). Steps: 1. Read each JSON into a list of dicts. 2. Compute total number of modes and states for all files. 3. Write one combined MCTDH file, adjusting offsets as needed to avoid collisions in modes, states, etc. Parameters: infiles (Union[str, List[str]]): One or more JSON file paths. outfile (str): Output MCTDH operator file path. """ # 1. Collect data from each JSON file data_blocks = _read_data_blocks(infiles) # 2. Compute total number of modes & states total_modes, total_states, master_units = _compute_totals(data_blocks) # 3) If dipole computations are requested if dipole: if not dipole_config: # Either create default or raise an error dipole_config = DipoleConfig() print("Dipole-dipole computations requested.") # 3. Write out the MCTDH operator file directly try: with open(outfile, "w") as fh: # 3A. Header _write_header_section(fh) # 3B. Parameter Section and calculate dipoles val_dipoles = _write_parameter_section(fh, data_blocks, master_units, dipoles=dipole, dipole_config=dipole_config) # 3C. Labels Section _write_labels_section(fh, data_blocks) # 3D. Hamiltonian Section _write_hamiltonian_section(fh, data_blocks, total_modes, total_states, dipole=dipole, dipole_config=dipole_config, values_dipoles=val_dipoles) logger.info(f"MCTDH operator file successfully written to: {outfile}") except IOError as e: logger.error(f"Error writing MCTDH file: {e}") raise