diff --git a/examples/sftda/04_sftda_soc.py b/examples/sftda/04_sftda_soc.py new file mode 120000 index 0000000..7725ed8 --- /dev/null +++ b/examples/sftda/04_sftda_soc.py @@ -0,0 +1 @@ +../soc/01_sftda_soc.py \ No newline at end of file diff --git a/examples/soc/01_sftda_soc.py b/examples/soc/01_sftda_soc.py new file mode 100644 index 0000000..74f66d9 --- /dev/null +++ b/examples/soc/01_sftda_soc.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python +# Copyright 2026 The NEST Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run an SF-TDA SOC calculation and print the complete SOC analysis.""" + +from pyscf import gto +from nest import sftda # noqa: F401 # Registers the SFTDA methods. + +mol = gto.M( + atom=""" + O 0.64372820 0.14077399 -0.04477253 + O -0.64862595 -0.12779073 -0.05445498 + H 1.16027512 -0.65947800 0.36730132 + H -1.12109306 0.55561188 0.42651873 + """, + basis="631g", + charge=0, + spin=2, + symmetry=False, +) + +# 1. High-spin scalar reference. +mf = mol.UKS(xc="SVWN").run() + +# 2. Spin-flip-down TDA states. SOC currently supports extype=1 only. +td = mf.SFTDA().set( + extype=1, + nstates=3, + collinear="mcol", + collinear_samples=50, +).run() + +# 3. Build and diagonalize the SOC Hamiltonian. +# +# The PySCF-style entry point below is recommended. run() calls kernel() and +# returns the SOC driver itself, with the numerical results stored as +# attributes. +soc_driver = td.SOC(soctype="SOMF").run() +# +# The equivalent explicit construction is: +# +# from nest import soc +# soc_driver = soc.sftda.SOC(td, soctype="SOMF") +# e, v = soc_driver.kernel() +# +# The first argument is the converged SF-TDA/SF-TDDFT object ``td``, not the +# SCF object ``mf``. Available soctype values are: +# +# "SOMF" one-electron SOC plus the two-electron SOMF contribution +# "Zeff" one-electron SOC with ORCA-style effective nuclear charges +# "1e" bare one-electron nuclear SOC +# "X2CAMF" X2CAMF SOC; requires the socutils package + +# 4. Print scalar states, SOC blocks, SOCCs, SOC energies, and compositions. +# verbose=4 enables the detailed block matrices and eigenstate compositions. +soc_driver.analyze(verbose=4) + +# The numerical results remain available for further processing: +# soc_driver.states spin-free input states +# soc_driver.state_slices mapping from spin-free states to Hamiltonian slices +# soc_driver.h_soc complex Hermitian SOC Hamiltonian, in Hartree +# soc_driver.e, .v SOC eigenvalues and eigenvectors diff --git a/src/nest/sftda/uhf_sf.py b/src/nest/sftda/uhf_sf.py index 07f4d5c..d949465 100644 --- a/src/nest/sftda/uhf_sf.py +++ b/src/nest/sftda/uhf_sf.py @@ -716,6 +716,11 @@ def NAC(self): return tduks_sf.NAC(self) nac_method = NAC + def SOC(self, soctype="SOMF"): + """Create an SOC driver from converged SF-TDA/SF-TDDFT results.""" + from nest.soc import sftda as sftda_soc + return sftda_soc.SOC(self, soctype=soctype) + analyze = analyze transition_dipole = transition_dipole oscillator_strength = oscillator_strength diff --git a/src/nest/soc/__init__.py b/src/nest/soc/__init__.py new file mode 100644 index 0000000..63ba203 --- /dev/null +++ b/src/nest/soc/__init__.py @@ -0,0 +1,24 @@ +#!/usr/bin/env python +# Copyright 2026 The NEST Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Author: Tai Wang & Codex +# + +"""Spin-orbit coupling drivers and AO integrals.""" + +from nest.soc.sftda import SOC as SFTDASOC +from nest.soc.soc import SOCBase, SpinFreeState + +__all__ = ['SOCBase', 'SpinFreeState', 'SFTDASOC'] diff --git a/src/nest/soc/sftda.py b/src/nest/soc/sftda.py new file mode 100644 index 0000000..bcf316b --- /dev/null +++ b/src/nest/soc/sftda.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python +# Copyright 2026 The NEST Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SOC driver for spin-flip TDA and TDDFT.""" + +import numpy as np +from pyscf import lib +from pyscf.lib import logger + +from nest.soc.soc import SOCBase, SpinFreeState, clebsch_gordan_rank1 + + +class SOC(SOCBase): + """Build SOC states from one converged SFTDA or SFTDDFT object.""" + + _keys = {'tdobj', 'sz'} + + def __init__(self, tdobj, soctype='SOMF'): + if tdobj.extype != 1: + raise ValueError('SFTDA SOC currently supports extype=1 only') + if getattr(tdobj, 'e', None) is None or getattr(tdobj, 'xy', None) is None: + raise ValueError('Run the SFTDA/SFTDDFT kernel before SOC') + assert np.isrealobj(tdobj._scf.mo_coeff), 'SFTDA SOC requires real MO coefficients' + + spin_square = np.asarray(tdobj.spin_square(), dtype=float) + states = [] + for root, energy in enumerate(tdobj.e): + s2 = float(spin_square[root]) + spin = round(-1 + np.sqrt(1 + 4 * s2)) / 2 + amplitude = tdobj.xy[root] + assert np.isrealobj(amplitude[0]), 'SFTDA SOC requires real X amplitudes' + if isinstance(amplitude[1], np.ndarray): + assert np.isrealobj(amplitude[1]), 'SFTDDFT SOC requires real Y amplitudes' + states.append(SpinFreeState( + source=tdobj, + root=root, + energy=float(energy), + spin=spin, + amplitude=amplitude, + label=f'SF state {root + 1}', + spin_square=s2, + )) + super().__init__(tdobj._scf, states, soctype=soctype) + self.tdobj = tdobj + self.sz = 0.5 * self._scf.mol.spin - 1 + + def reduced_transition_density(self, bra, ket): + nao = self._scf.mol.nao_nr() + if abs(bra.spin) < 1e-12 and abs(ket.spin) < 1e-12: + # A rank-one operator cannot couple two singlets. + return np.zeros((nao, nao), dtype=np.complex128) + + # Extract the reduced matrix element from the q=0 component at the + # spin projection represented by the spin-flip amplitudes. A zero + # coefficient here does not generally imply a zero reduced matrix + # element; it means that this q=0 component contains no information + # about it (for example, S=1, M_S=0 <-> S=1, M_S=0). + coefficient = clebsch_gordan_rank1(ket.spin, self.sz, 0, bra.spin, self.sz) + if abs(coefficient) < 1e-12: + logger.warn( + self, + 'The q=0 SF transition density cannot determine the reduced ' + 'SOC matrix element for S=%s <- S=%s at M_S=%s. Returning zero; ' + 'the affected states may be spin contaminated.', + bra.spin, ket.spin, self.sz, + ) + return np.zeros((nao, nao), dtype=np.complex128) + + mo_coeff = self._scf.mo_coeff + mo_occ = self._scf.mo_occ + occidxa = mo_occ[0] > 0 + occidxb = mo_occ[1] > 0 + viridxa = mo_occ[0] == 0 + viridxb = mo_occ[1] == 0 + orboa = mo_coeff[0][:, occidxa] + orbob = mo_coeff[1][:, occidxb] + orbva = mo_coeff[0][:, viridxa] + orbvb = mo_coeff[1][:, viridxb] + + mx = bra.amplitude[0] + nx = ket.amplitude[0] + gamma_oo_aa = -lib.einsum('ia,ja->ij', mx, nx) + gamma_aa = lib.einsum('uj,vi,ij->vu', orboa, orboa, gamma_oo_aa) + gamma_vv_bb = lib.einsum('ib,ia->ab', mx, nx) + gamma_bb = lib.einsum('ub,va,ab->vu', orbvb, orbvb, gamma_vv_bb) + + if isinstance(bra.amplitude[1], np.ndarray): + my = bra.amplitude[1] + ny = ket.amplitude[1] + gamma_oo_bb = -lib.einsum('ja,ia->ij', my, ny) + gamma_bb += lib.einsum('uj,vi,ij->vu', orbob, orbob, gamma_oo_bb) + gamma_vv_aa = lib.einsum('ia,ib->ab', my, ny) + gamma_aa += lib.einsum('ub,va,ab->vu', orbva, orbva, gamma_vv_aa) + + # The AO transformations above explicitly output ``gamma[v, u]``. + # Keep that physical transition-density order; SOCBase contracts it + # as z[u, v] * gamma[v, u]. + return (gamma_aa - gamma_bb) / np.sqrt(2) / coefficient + + +__all__ = ['SOC'] diff --git a/src/nest/soc/soc.py b/src/nest/soc/soc.py new file mode 100644 index 0000000..4fbde2a --- /dev/null +++ b/src/nest/soc/soc.py @@ -0,0 +1,248 @@ +#!/usr/bin/env python +# Copyright 2026 The NEST Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Common machinery for spin-orbit-coupled excited states.""" + +from dataclasses import dataclass +from math import sqrt +from typing import Any + +import numpy as np +from pyscf import lib +from pyscf.data.nist import HARTREE2EV, HARTREE2WAVENUMBER +from pyscf.lib import logger + +from nest.soc.soc_ao import get_ao_soc + + +@dataclass +class SpinFreeState: + """One scalar excited state before spin-orbit coupling is introduced.""" + + source: Any + root: int | None + energy: float + spin: float + amplitude: Any + label: str + spin_square: float | None = None + delta_s: int | None = None + + +def clebsch_gordan_rank1(j1, m1, q, j, m): + """Return ```` for the rank-one SOC tensor.""" + tol = 1e-12 + if q not in (-1, 0, 1) or abs(m - (m1 + q)) > tol: + return 0.0 + if j1 < 0 or j < 0 or abs(m1) > j1 + tol or abs(m) > j + tol: + return 0.0 + + if abs(j - (j1 + 1)) < tol: + if q == 1: + return sqrt((j1 + m1 + 1) * (j1 + m1 + 2) / (2 * (j1 + 1) * (2 * j1 + 1))) + if q == 0: + return sqrt((j1 - m1 + 1) * (j1 + m1 + 1) / ((j1 + 1) * (2 * j1 + 1))) + return sqrt((j1 - m1 + 1) * (j1 - m1 + 2) / (2 * (j1 + 1) * (2 * j1 + 1))) + + if abs(j - j1) < tol: + if j1 == 0: + return 0.0 + if q == 1: + return -sqrt((j1 - m1) * (j1 + m1 + 1) / (2 * j1 * (j1 + 1))) + if q == 0: + return m1 / sqrt(j1 * (j1 + 1)) + return sqrt((j1 + m1) * (j1 - m1 + 1) / (2 * j1 * (j1 + 1))) + + if abs(j - (j1 - 1)) < tol: + if j1 == 0: + return 0.0 + if q == 1: + return sqrt((j1 - m1) * (j1 - m1 - 1) / (2 * j1 * (2 * j1 + 1))) + if q == 0: + return -sqrt((j1 - m1) * (j1 + m1) / (j1 * (2 * j1 + 1))) + return sqrt((j1 + m1) * (j1 + m1 - 1) / (2 * j1 * (2 * j1 + 1))) + return 0.0 + + +class SOCBase(lib.StreamObject): + """Build and diagonalize a SOC Hamiltonian from scalar excited states.""" + + _keys = {'states', 'soctype', 'soc_ao', 'state_slices', 'h_soc', 'e', 'v'} + + def __init__(self, mf, states, soctype='SOMF'): + self._scf = mf + self.states = list(states) + self.soctype = soctype + self.verbose = getattr(mf, 'verbose', logger.NOTE) + self.stdout = getattr(mf, 'stdout', None) + + self.soc_ao = None + self.state_slices = None + self.h_soc = None + self.e = None + self.v = None + + @staticmethod + def _m_values(spin): + values = np.arange(-spin, spin + 0.5, 1.0) + values[abs(values) < 1e-12] = 0.0 + return values + + def reduced_transition_density(self, bra, ket): + r"""Return the AO reduced transition density for ``S_bra >= S_ket``. + \gamma_{\nu\mu} = + """ + raise NotImplementedError + + def _soc_block(self, bra, ket): + nbra = int(round(2 * bra.spin + 1)) + nket = int(round(2 * ket.spin + 1)) + if abs(bra.spin - ket.spin) > 1 + 1e-12: + return np.zeros((nbra, nket), dtype=np.complex128) + if bra.spin + 1e-12 < ket.spin: + return self._soc_block(ket, bra).conj().T + + density = self.reduced_transition_density(bra, ket) + # = sum_uv z_uv gamma_vu = Tr(z gamma). + # ``density[v, u]`` therefore contracts with ``soc_ao[x, u, v]``. + components = np.einsum('xuv,vu->x', self.soc_ao, density) + block = np.zeros((nbra, nket), dtype=np.complex128) + for row, m_bra in enumerate(self._m_values(bra.spin)): + for col, m_ket in enumerate(self._m_values(ket.spin)): + q = int(round(m_bra - m_ket)) + if abs(m_bra - m_ket - q) > 1e-12 or q not in (-1, 0, 1): + continue + coefficient = clebsch_gordan_rank1(ket.spin, m_ket, q, bra.spin, m_bra) + if q == 1: + block[row, col] = -coefficient * components[0] + elif q == 0: + block[row, col] = coefficient * components[1] + else: + block[row, col] = -coefficient * components[2] + return block + + def build_hamiltonian(self): + self.soc_ao = get_ao_soc(self._scf, self.soctype) + self.state_slices = [] + start = 0 + for state in self.states: + stop = start + len(self._m_values(state.spin)) + self.state_slices.append(slice(start, stop)) + start = stop + h_soc = np.zeros((start, start), dtype=np.complex128) + + # Diagonal first-order SOC blocks vanish for the real scalar states + # accepted by the method-specific drivers. + for state, state_slice in zip(self.states, self.state_slices): + dimension = state_slice.stop - state_slice.start + h_soc[state_slice, state_slice] = np.eye(dimension) * state.energy + + for bra_id in range(len(self.states)): + for ket_id in range(bra_id): + block = self._soc_block(self.states[bra_id], self.states[ket_id]) + bra_slice = self.state_slices[bra_id] + ket_slice = self.state_slices[ket_id] + h_soc[bra_slice, ket_slice] = block + h_soc[ket_slice, bra_slice] = block.conj().T + + error = np.max(abs(h_soc - h_soc.conj().T)) if h_soc.size else 0.0 + if error > 1e-10: + raise ValueError(f'SOC Hamiltonian is not Hermitian: max error {error:.3e}') + self.h_soc = h_soc + return h_soc + + def kernel(self): + self.build_hamiltonian() + self.e, self.v = np.linalg.eigh(self.h_soc) + return self.e, self.v + + def get_block(self, bra, ket): + if self.h_soc is None: + raise RuntimeError('Run kernel() or build_hamiltonian() first') + return self.h_soc[self.state_slices[bra], self.state_slices[ket]] + + @classmethod + def _format_soc_block(cls, block, bra, ket): + """Format one SOC block with explicit bra/ket ``M_S`` labels.""" + bra_m_values = cls._m_values(bra.spin) + ket_m_values = cls._m_values(ket.spin) + expected_shape = (len(bra_m_values), len(ket_m_values)) + if block.shape != expected_shape: + raise ValueError( + f'SOC block shape {block.shape} does not match the expected ' + f'bra/ket dimensions {expected_shape}' + ) + + label_width = 18 + column_width = 27 + header = ' ' * label_width + for m_ket in ket_m_values: + header += f'ket M_S={m_ket:5.1f}'.center(column_width) + + lines = ['SOC matrix elements (cm^-1):', header, '-' * len(header)] + for row, m_bra in enumerate(bra_m_values): + line = f'bra M_S={m_bra:5.1f}'.ljust(label_width) + for value in block[row]: + line += f'({value.real:10.6f},{value.imag:10.6f})'.center(column_width) + lines.append(line) + return lines + + def analyze(self, verbose=None): + if self.h_soc is None or self.e is None: + self.kernel() + log = logger.new_logger(self, verbose) + log.note('SOC scalar states') + for state_id, state in enumerate(self.states): + message = 'State %d: %s S=%s E=%.8f Eh' % ( + state_id + 1, state.label, state.spin, state.energy, + ) + if state.spin_square is not None: + message += ' =%.6f' % state.spin_square + log.note(message) + + for bra_id in range(len(self.states)): + for ket_id in range(bra_id): + block = self.get_block(bra_id, ket_id) + bra = self.states[bra_id] + ket = self.states[ket_id] + log.note( + 'SOC block: state %d (%s, S=%s) <- state %d (%s, S=%s); ' + 'SOCC = %.6f cm^-1', + bra_id + 1, bra.label, bra.spin, + ket_id + 1, ket.label, ket.spin, + np.linalg.norm(block) * HARTREE2WAVENUMBER, + ) + for line in self._format_soc_block( + block * HARTREE2WAVENUMBER, bra, ket, + ): + log.info('%s', line) + + origin = self.e.min() if self.e.size else 0.0 + log.note('Spin-orbit-coupled eigenstates') + for state_id, energy in enumerate(self.e): + log.note('State %d: Delta E = %.6f cm^-1 (%.8f eV)', state_id + 1, + (energy - origin).real * HARTREE2WAVENUMBER, + (energy - origin).real * HARTREE2EV) + weights = abs(self.v[:, state_id]) ** 2 + log.info(' Spin-free state composition:') + for state, state_slice in zip(self.states, self.state_slices): + probability = weights[state_slice].sum() + if probability > 0.01: + log.info(' %5.1f%% from %s (S=%s)', + probability * 100, state.label, state.spin) + return self + + +__all__ = ['SpinFreeState', 'SOCBase', 'clebsch_gordan_rank1'] diff --git a/src/nest/soc/soc_ao.py b/src/nest/soc/soc_ao.py new file mode 100644 index 0000000..92e4bcc --- /dev/null +++ b/src/nest/soc/soc_ao.py @@ -0,0 +1,262 @@ +#!/usr/bin/env python +# Copyright 2026 The NEST Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SOC Hamiltonian in AO basis.""" + +import numpy as np +from pyscf.scf.jk import get_jk +from pyscf.data.nist import LIGHT_SPEED + +X2CAMF_XRESP = True + +def sozeff(atom, zeff_type="one"): + """ + Calculate effective nuclear charge for given atomic number + copied from: https://github.com/masaya0222/PyGraSO/blob/main/pygraso/calc_ao_element.py + Author: Masaya Hagai + """ + assert zeff_type in ["one", "orca", "pysoc"], f"{zeff_type=} is not valid" + neval = { + 1: 1, + 2: 2, + 3: 1, + 4: 2, + 5: 3, + 6: 4, + 7: 5, + 8: 6, + 9: 7, + 10: 8, + 11: 1, + 12: 2, + 13: 3, + 14: 4, + 15: 5, + 16: 6, + 17: 7, + 18: 8, + 19: 1, + 20: 2, + 21: 3, + 22: 4, + 23: 5, + 24: 6, + 25: 7, + 26: 8, + 27: 9, + 28: 10, + 29: 11, + 30: 12, + 31: 3, + 32: 4, + 33: 5, + 34: 6, + 35: 7, + 36: 8, + 37: 1, + 38: 2, + 39: 3, + 40: 4, + 41: 5, + 42: 6, + 43: 7, + 44: 8, + 45: 9, + 46: 10, + 47: 11, + 48: 12, + 49: 3, + 50: 4, + 51: 5, + 52: 6, + 53: 7, + 54: 8, + } + if zeff_type == "one": + return atom + + if zeff_type == "pysoc": + if atom == 1: + return 1.0 + elif atom == 2: + return 2.0 + elif 3 <= atom <= 10: + return (0.2517 + 0.0626 * neval[atom]) * atom + elif 11 <= atom <= 18: + return (0.7213 + 0.0144 * neval[atom]) * atom + elif (19 <= atom <= 20) or (31 <= atom <= 36): + return (0.8791 + 0.0039 * neval[atom]) * atom + elif (37 <= atom <= 38) or (49 <= atom <= 54): + return (0.9228 + 0.0017 * neval[atom]) * atom + elif atom == 26: + return 0.583289 * atom + elif atom == 30: + return 330.0 + elif 21 <= atom <= 30: + return atom * (0.385 + 0.025 * (neval[atom] - 2)) + elif 39 <= atom <= 48: + return atom * (4.680 + 0.060 * (neval[atom] - 2)) + elif atom == 72: + return 1025.28 + elif atom == 73: + return 1049.74 + elif atom == 74: + return 1074.48 + elif atom == 75: + return 1099.5 + elif atom == 76: + return 1124.8 + elif atom == 77: + return 1150.38 + elif atom == 78: + return 1176.24 + elif atom == 79: + return 1202.38 + elif atom == 80: + return 1228.8 + else: + raise ValueError(f"SOZEFF is not available for atomic number {atom}") + if zeff_type == "orca": + if atom == 1: + return 1.0 + elif atom == 2: + return 2.0 + elif 3 <= atom < 10: + return (0.4 + 0.05 * neval[atom]) * atom + elif 11 <= atom <= 18: + return (0.925 - 0.0125 * neval[atom]) * atom + elif 32 <= atom <= 35: # Verified from orca output file + if atom == 32: + return 32.32 + elif atom == 33: + return 31.68 + elif atom == 34: + return 30.94 + elif atom == 35: + return 30.10 + else: + raise ValueError(f"SOZEFF is not available for atomic number {atom}") + +def get_ao_soc_1e(mol, zeff_type='one'): + ''' + The one-body part of Hsoc operator with (effective) nuclear charge. + ''' + zeff_list = [sozeff(mol.atom_charge(i), zeff_type=zeff_type) for i in range(mol.natm)] + ao_soc = np.zeros((3, mol.nao_nr(), mol.nao_nr()), dtype=np.complex128) + for k in range(mol.natm): + mol.set_rinv_orig(mol.atom_coord(k)) + ao_soc += (-1.0j) * zeff_list[k] * mol.intor('int1e_prinvxp') + ao_soc /= (2.0 * LIGHT_SPEED**2) + ao_soc_1 = -0.5 * (ao_soc[0] + 1j * ao_soc[1]) + ao_soc_0 = np.sqrt(0.5) * ao_soc[2] + ao_soc_m1 = 0.5 * (ao_soc[0] - 1j * ao_soc[1]) + return np.array([ao_soc_m1, ao_soc_0, ao_soc_1]) + +def get_ao_soc_x2camf(mol): + import resource + try: + soft, hard = resource.getrlimit(resource.RLIMIT_STACK) + target = resource.RLIM_INFINITY if hard == resource.RLIM_INFINITY else hard + resource.setrlimit(resource.RLIMIT_STACK, (target, hard)) + except Exception as e: + print(f"Warning: failed to increase stack size: {e}") + try: + from socutils.somf import somf_pt + except ImportError: + raise ImportError("Please install socutils package to use X2CAMF SOC integrals." \ + "https://github.com/wtpeter/socutils") + ao_soc = 2j * somf_pt.get_psoc_x2camf(mol, xresp=X2CAMF_XRESP) + ao_soc_1 = -0.5 * (ao_soc[0] + 1j * ao_soc[1]) + ao_soc_0 = np.sqrt(0.5) * ao_soc[2] + ao_soc_m1 = 0.5 * (ao_soc[0] - 1j * ao_soc[1]) + return np.array([ao_soc_m1, ao_soc_0, ao_soc_1]) + +def get_ao_soc_2e_somf(mf): + ''' + The two-electron part of the Hsoc under the SOMF approximation + using direct SCF for UKS density matrix. + ''' + mol = mf.mol + dm = mf.make_rdm1() + + if dm.ndim == 2: + dmaa = dmbb = 0.5 * dm + else: + dmaa, dmbb = dm + + dm_list = [dmaa, dmaa, dmaa, dmbb, dmbb, dmbb] + scripts = [ + 'ijkl,lk->ij', # J for dmaa + 'ijkl,jk->il', # K1 for dmaa + 'ijkl,li->kj', # K2 for dmaa + 'ijkl,lk->ij', # J for dmbb + 'ijkl,jk->il', # K1 for dmbb + 'ijkl,li->kj' # K2 for dmbb + ] + v_matrices = get_jk(mol, dm_list, scripts=scripts, intor='int2e_p1vxp1', comp=3, aosym='a4ij') + vj_aa, vk1_aa, vk2_aa = v_matrices[0:3] + vj_bb, vk1_bb, vk2_bb = v_matrices[3:6] + + v_cart_1 = (vj_aa - vk1_aa - 2 * vk2_aa) + (vj_bb - 2 * vk1_bb - vk2_bb) + v_cart_0 = (vj_aa + vj_bb) - 1.5 * (vk1_aa + vk1_bb) - 1.5 * (vk2_aa + vk2_bb) + v_cart_m1 = (vj_aa - 2 * vk1_aa - vk2_aa) + (vj_bb - vk1_bb - 2 * vk2_bb) + + def to_spherical(v_cart_xyz, component): + # v_cart_xyz shape is (3, nao, nao) + vx = v_cart_xyz[0] + vy = v_cart_xyz[1] + vz = v_cart_xyz[2] + if component == 1: + return -0.5 * (vx + 1j * vy) + elif component == 0: + return np.sqrt(0.5) * vz + elif component == -1: + return 0.5 * (vx - 1j * vy) + + prefactor = 1j / (2.0 * LIGHT_SPEED**2) + soc_somf_1 = to_spherical(v_cart_1, 1) * prefactor + soc_somf_0 = to_spherical(v_cart_0, 0) * prefactor + soc_somf_m1 = to_spherical(v_cart_m1, -1) * prefactor + soc_somf_1, soc_somf_m1 = ( + 0.5 * (soc_somf_1 + soc_somf_m1.conj()), + 0.5 * (soc_somf_m1 + soc_somf_1.conj()), + ) + return np.array([soc_somf_m1, soc_somf_0, soc_somf_1]) + + +def _symmetrize_ao_soc(soc_ao): + """Enforce the Hermiticity relations of a rank-one spherical tensor.""" + soc_m1, soc_0, soc_1 = soc_ao + soc_0 = 0.5 * (soc_0 + soc_0.conj().T) + soc_m1, soc_1 = ( + 0.5 * (soc_m1 - soc_1.conj().T), + 0.5 * (soc_1 - soc_m1.conj().T), + ) + return np.array([soc_m1, soc_0, soc_1]) + +def get_ao_soc(mf, soctype): + mol = mf.mol + if soctype == 'SOMF': + soc_ao = get_ao_soc_1e(mol, zeff_type='one') + soc_ao += get_ao_soc_2e_somf(mf) + elif soctype == 'Zeff': + soc_ao = get_ao_soc_1e(mol, zeff_type='orca') + elif soctype == '1e': + soc_ao = get_ao_soc_1e(mol, zeff_type='one') + elif soctype == 'X2CAMF': + soc_ao = get_ao_soc_x2camf(mol) + else: + raise ValueError(f'soctype={soctype} is not supported.') + return _symmetrize_ao_soc(soc_ao) diff --git a/src/nest/soc/tests/test_sftda_soc.py b/src/nest/soc/tests/test_sftda_soc.py new file mode 100644 index 0000000..94259c9 --- /dev/null +++ b/src/nest/soc/tests/test_sftda_soc.py @@ -0,0 +1,107 @@ +# Copyright 2026 The NEST Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. + +import unittest +import numpy as np +from pyscf import gto +from pyscf.data.nist import HARTREE2WAVENUMBER +from nest import sftda +from nest.soc.soc import clebsch_gordan_rank1 + + +def assert_allclose_up_to_sign(testcase, actual, desired, atol): + try: + np.testing.assert_allclose(actual, desired, atol=atol, rtol=0) + except AssertionError as error: + np.testing.assert_allclose(actual, -desired, atol=atol, rtol=0, err_msg=str(error)) + + +class KnownValues(unittest.TestCase): + @classmethod + def setUpClass(cls): + mol = gto.Mole() + mol.verbose = 0 + mol.output = '/dev/null' + mol.atom = ''' + O 0.64372820 0.14077399 -0.04477253 + O -0.64862595 -0.12779073 -0.05445498 + H 1.16027512 -0.65947800 0.36730132 + H -1.12109306 0.55561188 0.42651873 + ''' + mol.charge = 0 + mol.spin = 2 + mol.basis = '631g' + cls.mol = mol.build() + + @classmethod + def tearDownClass(cls): + cls.mol.stdout.close() + + def test_rank_one_clebsch_gordan(self): + self.assertAlmostEqual(clebsch_gordan_rank1(0, 0, 1, 1, 1), 1.0) + self.assertAlmostEqual(clebsch_gordan_rank1(1, 1, 0, 1, 1), 2 ** -0.5) + self.assertAlmostEqual(clebsch_gordan_rank1(1, 0, 1, 1, 1), -2 ** -0.5) + self.assertAlmostEqual(clebsch_gordan_rank1(1, 0, 0, 0, 0), -3 ** -0.5) + + def test_roks_sftda_soc(self): + mf = self.mol.ROKS(xc='SVWN').run() + td = sftda.TDA_SF(mf).set( + extype=1, collinear='mcol', collinear_samples=50, nstates=3, + ).run() + driver = td.SOC(soctype='SOMF') + driver.kernel() + + self.assertTrue(mf.converged) + self.assertTrue(np.all(td.converged)) + self.assertAlmostEqual(mf.e_tot, -150.18173594947896, delta=1e-9) + np.testing.assert_allclose(td.e, [ + -0.2104295981711506, -0.0007174487394460, 0.0251523165536107, + ], atol=1e-8, rtol=0) + np.testing.assert_allclose(td.spin_square(), [ + 0.0010848962999618905, 1.9999490812871423, 0.031289468589663194, + ], atol=1e-8, rtol=0) + self.assertEqual(driver.h_soc.shape, (5, 5)) + np.testing.assert_allclose(driver.h_soc, driver.h_soc.conj().T, atol=1e-12) + np.testing.assert_allclose((driver.e - driver.e.min()).real * HARTREE2WAVENUMBER, [ + 0.0, 46026.50192719676, 46026.502933180134, 46026.50860068575, 51704.26176254972, + ], atol=1e-5, rtol=0) + assert_allclose_up_to_sign(self, driver.get_block(1, 0) * HARTREE2WAVENUMBER, np.array([ + [0.4674642361078794 - 6.685128880147436j], + [0.0 - 14.141141940607179j], + [0.4674642361078794 + 6.685128880147436j], + ]), 1e-8) + + def test_uks_sftda_soc(self): + mf = self.mol.UKS(xc='SVWN').run() + td = sftda.TDA_SF(mf).set( + extype=1, collinear='mcol', collinear_samples=50, nstates=3, + ).run() + driver = td.SOC(soctype='SOMF') + driver.kernel() + + self.assertTrue(mf.converged) + self.assertTrue(np.all(td.converged)) + self.assertAlmostEqual(mf.e_tot, -150.18252681880003, delta=1e-9) + np.testing.assert_allclose(td.e, [ + -0.2087681123003969, 0.0008054142507056, 0.0266315014304553, + ], atol=1e-8, rtol=0) + np.testing.assert_allclose(td.spin_square(), [ + 0.0026539116310360, 2.0039691808012283, 0.0372918876839510, + ], atol=1e-8, rtol=0) + self.assertEqual(driver.h_soc.shape, (5, 5)) + np.testing.assert_allclose(driver.h_soc, driver.h_soc.conj().T, atol=1e-12) + np.testing.assert_allclose((driver.e - driver.e.min()).real * HARTREE2WAVENUMBER, [ + 0.0, 45996.07760059907, 45996.07872056033, 45996.08436902181, 51664.25143819543, + ], atol=1e-5, rtol=0) + assert_allclose_up_to_sign(self, driver.get_block(1, 0) * HARTREE2WAVENUMBER, np.array([ + [-0.4695042701337793 + 6.699202067333119j], + [0.0 + 14.109500425297364j], + [-0.4695042701337793 - 6.699202067333119j], + ]), 1e-8) + + +if __name__ == '__main__': + print('Full SOC tests for spin-flip TDA') + unittest.main() diff --git a/src/nest/soc/tests/test_sftddft_soc.py b/src/nest/soc/tests/test_sftddft_soc.py new file mode 100644 index 0000000..3af3f96 --- /dev/null +++ b/src/nest/soc/tests/test_sftddft_soc.py @@ -0,0 +1,100 @@ +# Copyright 2026 The NEST Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. + +import unittest +import numpy as np +from pyscf import gto +from pyscf.data.nist import HARTREE2WAVENUMBER +from nest import sftda + + +def assert_allclose_up_to_sign(testcase, actual, desired, atol): + try: + np.testing.assert_allclose(actual, desired, atol=atol, rtol=0) + except AssertionError as error: + np.testing.assert_allclose(actual, -desired, atol=atol, rtol=0, err_msg=str(error)) + + +class KnownValues(unittest.TestCase): + @classmethod + def setUpClass(cls): + mol = gto.Mole() + mol.verbose = 0 + mol.output = '/dev/null' + mol.atom = ''' + O 0.64372820 0.14077399 -0.04477253 + O -0.64862595 -0.12779073 -0.05445498 + H 1.16027512 -0.65947800 0.36730132 + H -1.12109306 0.55561188 0.42651873 + ''' + mol.charge = 0 + mol.spin = 2 + mol.basis = '631g' + cls.mol = mol.build() + + @classmethod + def tearDownClass(cls): + cls.mol.stdout.close() + + def test_roks_sftddft_soc(self): + mf = self.mol.ROKS(xc='SVWN').run() + td = sftda.TDDFT_SF(mf).set( + extype=1, collinear='mcol', collinear_samples=50, nstates=3, + ).run() + driver = td.SOC(soctype='SOMF') + driver.kernel() + + self.assertTrue(mf.converged) + self.assertTrue(np.all(td.converged)) + self.assertAlmostEqual(mf.e_tot, -150.18173594947856, delta=1e-9) + np.testing.assert_allclose(td.e, [ + -0.2107471654542317, -0.0015844442063221, 0.0245430214421581, + ], atol=1e-8, rtol=0) + np.testing.assert_allclose(td.spin_square(), [ + 0.0011758949708333688, 2.0012261115407766, 0.034155774522082627, + ], atol=1e-8, rtol=0) + self.assertEqual(driver.h_soc.shape, (5, 5)) + np.testing.assert_allclose(driver.h_soc, driver.h_soc.conj().T, atol=1e-12) + np.testing.assert_allclose((driver.e - driver.e.min()).real * HARTREE2WAVENUMBER, [ + 0.0, 45905.916322187106, 45905.91746253988, 45905.92316008696, 51640.23516406501, + ], atol=1e-5, rtol=0) + assert_allclose_up_to_sign(self, driver.get_block(1, 0) * HARTREE2WAVENUMBER, np.array([ + [0.4626543873150336 - 6.686744761741945j], + [0.0 - 14.2382866761041j], + [0.4626543873150336 + 6.686744761741945j], + ]), 1e-8) + + def test_uks_sftddft_soc(self): + mf = self.mol.UKS(xc='SVWN').run() + td = sftda.TDDFT_SF(mf).set( + extype=1, collinear='mcol', collinear_samples=50, nstates=3, + ).run() + driver = td.SOC(soctype='SOMF') + driver.kernel() + + self.assertTrue(mf.converged) + self.assertTrue(np.all(td.converged)) + self.assertAlmostEqual(mf.e_tot, -150.18252681880003, delta=1e-9) + np.testing.assert_allclose(td.e, [ + -0.20907621837286508, 0.0000011484871538744751, 0.026028246454667784, + ], atol=1e-8, rtol=0) + np.testing.assert_allclose(td.spin_square(), [ + 0.0027459512964301, 2.0084594505481292, 0.0400313245672246, + ], atol=1e-8, rtol=0) + self.assertEqual(driver.h_soc.shape, (5, 5)) + np.testing.assert_allclose(driver.h_soc, driver.h_soc.conj().T, atol=1e-12) + np.testing.assert_allclose((driver.e - driver.e.min()).real * HARTREE2WAVENUMBER, [ + 0.0, 45887.18304976328, 45887.18432227726, 45887.189989714265, 51599.47400884429, + ], atol=1e-5, rtol=0) + assert_allclose_up_to_sign(self, driver.get_block(1, 0) * HARTREE2WAVENUMBER, np.array([ + [0.4709967390252593 - 6.694129899974075j], + [0.0 - 14.1984224305132j], + [0.4709967390252593 + 6.694129899974075j], + ]), 1e-8) + + +if __name__ == '__main__': + print('Full SOC tests for spin-flip TDDFT') + unittest.main() diff --git a/src/nest/soc/tests/test_soc_ao.py b/src/nest/soc/tests/test_soc_ao.py new file mode 100644 index 0000000..41a064c --- /dev/null +++ b/src/nest/soc/tests/test_soc_ao.py @@ -0,0 +1,83 @@ +# Copyright 2026 The NEST Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest +from pyscf import gto, lib +from nest.soc import soc_ao + + +def fp(mat): + return lib.fp(mat) + +class KnownValues(unittest.TestCase): + @classmethod + def setUpClass(cls): + mol = gto.Mole() + mol.verbose = 0 + mol.output = '/dev/null' + mol.atom = ''' + O 0.64372820 0.14077399 -0.04477253 + O -0.64862595 -0.12779073 -0.05445498 + H 1.16027512 -0.65947800 0.36730132 + H -1.12109306 0.55561188 0.42651873 + ''' + mol.charge = 0 + mol.spin = 2 + mol.basis = '631g' + cls.mol = mol.build() + + @classmethod + def tearDownClass(cls): + cls.mol.stdout.close() + + def test_1e_soc_ao(self): + mf = self.mol.ROKS(xc='HF').run() + self.assertTrue(mf.converged) + + ref = -0.0035217717150048196 - 0.002071753516350664j + ao_soc = soc_ao.get_ao_soc(mf, '1e') + self.assertEqual(ao_soc.shape, (3, self.mol.nao_nr(), self.mol.nao_nr())) + self.assertAlmostEqual(abs(fp(ao_soc) - ref), 0, delta=1e-11) + + def test_zeff_soc_ao(self): + mf = self.mol.ROKS(xc='HF').run() + self.assertTrue(mf.converged) + + ref = -0.002466043752900166 - 0.001450202630169735j + ao_soc = soc_ao.get_ao_soc(mf, 'Zeff') + self.assertEqual(ao_soc.shape, (3, self.mol.nao_nr(), self.mol.nao_nr())) + self.assertAlmostEqual(abs(fp(ao_soc) - ref), 0, delta=1e-11) + + def test_somf_soc_ao(self): + mf = self.mol.ROKS(xc='HF').run() + self.assertTrue(mf.converged) + + ref = -0.002331075720748517 - 0.0013735097397886604j + ao_soc = soc_ao.get_ao_soc(mf, 'SOMF') + self.assertEqual(ao_soc.shape, (3, self.mol.nao_nr(), self.mol.nao_nr())) + self.assertAlmostEqual(abs(fp(ao_soc) - ref), 0, delta=1e-9) + + def test_somf_soc_ao_uks(self): + mf = self.mol.UKS(xc='HF').newton().run() + self.assertTrue(mf.converged) + + ref = -0.0023423479533036867 - 0.0013759118729275677j + ao_soc = soc_ao.get_ao_soc(mf, 'SOMF') + self.assertEqual(ao_soc.shape, (3, self.mol.nao_nr(), self.mol.nao_nr())) + self.assertAlmostEqual(abs(fp(ao_soc) - ref), 0, delta=1e-9) + + +if __name__ == '__main__': + print('Full tests for AO spin-orbit integrals') + unittest.main()