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