diff --git a/pysisyphus/drivers/spectrum_ctnum.py b/pysisyphus/drivers/spectrum_ctnum.py index 92edde75c..a6a26936e 100644 --- a/pysisyphus/drivers/spectrum_ctnum.py +++ b/pysisyphus/drivers/spectrum_ctnum.py @@ -29,7 +29,8 @@ def get_frags_noh(atoms, coords, bond_inds): return frags -def plot_geometry_graph(ax, atoms, coords3d, frags, bond_inds, scale=12.5): +def plot_geometry_graph(ax, atoms, coords3d, frags, bond_inds, scale=12.5, filter_atoms={"h", }): + atoms = [atom.lower() for atom in atoms] node_color = list() node_size = list() pos0 = dict() @@ -39,13 +40,15 @@ def plot_geometry_graph(ax, atoms, coords3d, frags, bond_inds, scale=12.5): drop_ind = pos_variances.argmin() keep_dims = [i for i in range(3) if i != drop_ind] for i, atom in enumerate(atoms): - atom = atom.lower() node_color.append(CPK_RGB[atom]) node_size.append(COVALENT_RADII[atom]) - if atom == "h": + if atom in filter_atoms: continue pos0[i] = coords3d[i, keep_dims] G.add_node(i, atom=atom) + pos0_max = coords3d[:, keep_dims].max(axis=0) + pos0_min = coords3d[:, keep_dims].min(axis=0) + pos0_extents = pos0_max - pos0_min node_color = np.array(node_color) node_size = np.array(node_size) node_size /= node_size.min() @@ -54,15 +57,28 @@ def plot_geometry_graph(ax, atoms, coords3d, frags, bond_inds, scale=12.5): for from_, to_ in bond_inds: from_atom = atoms[from_] to_atom = atoms[to_] - if from_atom == "h" or to_atom == "h": + if from_atom in filter_atoms or to_atom in filter_atoms: continue bond_length = np.linalg.norm(coords3d[from_] - coords3d[to_]) G.add_edge(from_, to_, weight=bond_length) labels = { - i: f"{atom.capitalize()}{i}" for i, atom in enumerate(atoms) if atom != "h" + i: f"{atom.capitalize()}{i}" for i, atom in enumerate(atoms) if atom not in filter_atoms } - pos = nx.kamada_kawai_layout(G, scale=scale, pos=pos0) + try: + pos = nx.nx_agraph.graphviz_layout(G, prog="neato") + pos_arr = np.array(list(pos.values())) + pos_max = pos_arr.max(axis=0) + pos_min = pos_arr.min(axis=0) + pos_extents = pos_max - pos_min + pos_scale = pos0_extents / pos_extents + pos_arr = pos_arr * pos_scale + for i, key in enumerate(pos.keys()): + pos[key] = pos_arr[i] + except ImportError: + print("Falling back to networkx-layout, as (py)graphviz is not available.") + # Kamda-Kawai can also lead to overlapping nodes + pos = nx.kamada_kawai_layout(G, scale=scale, pos=pos0) """ nx.draw( G, pos=pos, ax=ax, labels=labels, node_color=node_color, node_size=node_size @@ -74,7 +90,16 @@ def plot_geometry_graph(ax, atoms, coords3d, frags, bond_inds, scale=12.5): # nodes different, depending on whether they loose or receive electron density. # Now arrows are used and drawing each fragment separately is actually not required. # But it doesn't hurt either to keep it like this. - for fnodes in frags: + frags_filtered = list() + for frag in frags: + cur_frag = list() + for atom_ind in frag: + if atoms[atom_ind] in filter_atoms: + continue + cur_frag.append(atom_ind) + frags_filtered.append(cur_frag) + + for fnodes in frags_filtered: fnode_color = node_color[fnodes] fnode_size = node_size[fnodes] fpos = sum([pos[n] for n in fnodes]) / len(fnodes) @@ -204,6 +229,9 @@ def render_exc(ind, thresh=0.1): def get_annotate_state(ax, exc_ens, foscs): exc_ens_nm = _1OVER_AU2NM / exc_ens exc_ens_eV = exc_ens * AU2EV + ylim = ax.get_ylim() + ymax = ylim[1] + lower_third = (ymax - ylim[0]) / 3. def annotate_state(ind): exc_en_nm = exc_ens_nm[ind] @@ -211,7 +239,8 @@ def annotate_state(ind): fosc = foscs[ind] text = f"{exc_en_nm:.2f} nm, {exc_en_eV:.2f} eV\nf={fosc:.5f}" xy = (exc_en_nm, fosc) - xytext = (exc_en_nm + 15.0, fosc + 0.125) + ytext = min(ymax, fosc + max(fosc * 0.1, lower_third)) + xytext = (exc_en_nm + 15.0, ytext)#fosc + max(fosc * 0.1, lower_third)) annot = ax.annotate( text, xy, diff --git a/pysisyphus/intcoords/setup.py b/pysisyphus/intcoords/setup.py index 8637110a2..0b9c98684 100644 --- a/pysisyphus/intcoords/setup.py +++ b/pysisyphus/intcoords/setup.py @@ -72,12 +72,16 @@ def get_fragments( bond_inds=None, bond_factor=BOND_FACTOR, ignore_atom_inds=None, + ignore_bonds=None, with_unconnected_atoms=False, ): """This misses unconnected single atoms w/ 'with_unconnected_atoms=False'!""" coords3d = coords.reshape(-1, 3) if ignore_atom_inds is None: ignore_atom_inds = list() + if ignore_bonds is None: + ignore_bonds = set() + ignore_bonds = set([frozenset(ib) for ib in ignore_bonds]) ignore_atom_inds = set(ignore_atom_inds) if bond_inds is None: @@ -85,7 +89,8 @@ def get_fragments( bond_inds = get_bond_sets(atoms, coords3d, bond_factor=bond_factor) bond_ind_sets = [frozenset(bi) for bi in bond_inds] - bond_ind_sets = [bi for bi in bond_ind_sets if not bi & ignore_atom_inds] + bond_ind_sets = [bi for bi in bond_ind_sets + if (not bi & ignore_atom_inds) and (bi not in ignore_bonds)] fragments = merge_sets(bond_ind_sets) # Also add single-atom fragments for unconnected atoms that don't participate diff --git a/scripts/orca_ctnums.py b/scripts/orca_ctnums.py index 1083e83e7..7b7504942 100755 --- a/scripts/orca_ctnums.py +++ b/scripts/orca_ctnums.py @@ -2,8 +2,10 @@ import argparse +import itertools as it from pathlib import Path import sys +from typing import Optional import matplotlib.pyplot as plt import numpy as np @@ -15,7 +17,10 @@ from pysisyphus.wavefunction.excited_states import norm_ci_coeffs, ct_numbers_for_states -def load_data(base_name: Path): +def load_data(base_name: Path, triplets: bool = False, ignore_bonds: Optional[list[int]] = None): + if ignore_bonds is None: + ignore_bonds = list() + dump_fn = base_name.with_suffix(".npz") if dump_fn.exists(): @@ -27,7 +32,7 @@ def load_data(base_name: Path): wf_fn = base_name.with_suffix(".bson") # Drop GS, only keep excitation energies - all_ens = parse_orca_all_energies(log_fn, do_tddft=True) + all_ens = parse_orca_all_energies(log_fn, triplets=triplets, do_tddft=True) exc_ens = all_ens[1:] - all_ens[0] print(f"Parsed excitation energies from '{log_fn}'.") @@ -36,7 +41,7 @@ def load_data(base_name: Path): wf = Wavefunction.from_file(wf_fn) print(f"Loaded wavefunction from '{wf_fn}'.") - frags = get_fragments(wf.atoms, wf.coords, with_unconnected_atoms=True) + frags = get_fragments(wf.atoms, wf.coords, ignore_bonds=ignore_bonds, with_unconnected_atoms=True) # Convert from list of frozensets to list of lists frags = list(map(list, frags)) @@ -52,6 +57,12 @@ def load_data(base_name: Path): foscs = wf.oscillator_strength(exc_ens, tdms) print(f"Calculated oscillator strengths.") + homogeneous_frags = np.zeros((len(wf.atoms), 2), dtype=int) + i = 0 + for j, frag in enumerate(frags): + for atom in frag: + homogeneous_frags[i] = j, atom + i += 1 data = dict( atoms=wf.atoms, coords=wf.coords, @@ -64,7 +75,7 @@ def load_data(base_name: Path): tdms=tdms, foscs=foscs, ct_numbers=ct_numbers, - frags=frags, + frags=homogeneous_frags, ) np.savez(dump_fn, **data) return data @@ -74,6 +85,8 @@ def parse_args(args): parser = argparse.ArgumentParser() parser.add_argument("base_name") + parser.add_argument("--triplets", action="store_true") + parser.add_argument("--bonds", nargs="+", type=int) return parser.parse_args(args) @@ -82,14 +95,32 @@ def run(): args = parse_args(sys.argv[1:]) base_name = Path(args.base_name) - data = load_data(base_name) + triplets = args.triplets + bonds = args.bonds + if bonds is None: + bonds = list() + + assert len(bonds) % 2 == 0, "Length of bonds must be a multiple of 2!" + nbonds = len(bonds) // 2 + ignore_bonds = list() + for i in range(nbonds): + ignore_bonds.append(bonds[2*i:2*(i+1)]) + + data = load_data(base_name, triplets, ignore_bonds) atoms = data["atoms"] coords = data["coords"] exc_ens = data["exc_ens"] foscs = data["foscs"] ct_numbers = data["ct_numbers"] - frags = data["frags"] + homogenous_frags = data["frags"] + + frags = [] + for key, _frag in it.groupby(homogenous_frags, key=lambda frag_atom_ind: frag_atom_ind[0]): + cur_frag = [] + for _, atom_ind in _frag: + cur_frag.append(atom_ind) + frags.append(cur_frag) fig = ct_number_plot(atoms, coords, exc_ens, foscs, ct_numbers, frags=frags) plt.show()