75 lines
2.1 KiB
Python
75 lines
2.1 KiB
Python
import numpy as np
|
|
from dataclasses import dataclass
|
|
from typing import List
|
|
|
|
ATOMIC_NUMBERS = {
|
|
'H': 1, 'He': 2, 'Li': 3, 'Be': 4, 'B': 5,
|
|
'C': 6, 'N': 7, 'O': 8, 'F': 9, 'Ne': 10,
|
|
# Soft pseudopotential variants (same atomic number)
|
|
'H_soft': 1,
|
|
}
|
|
ANGSTROM_TO_BOHR = 1.8897259886
|
|
|
|
|
|
@dataclass
|
|
class Atom:
|
|
symbol: str
|
|
coords: np.ndarray # Bohr
|
|
|
|
@property
|
|
def Z(self) -> int:
|
|
return ATOMIC_NUMBERS[self.symbol]
|
|
|
|
|
|
class Molecule:
|
|
def __init__(self, atoms: List[Atom], charge: int = 0, spin: int = 0):
|
|
self.atoms = atoms
|
|
self.charge = charge
|
|
self.spin = spin
|
|
|
|
@property
|
|
def n_electrons(self) -> int:
|
|
return sum(a.Z for a in self.atoms) - self.charge
|
|
|
|
@property
|
|
def n_alpha(self) -> int:
|
|
return (self.n_electrons + self.spin) // 2
|
|
|
|
@property
|
|
def n_beta(self) -> int:
|
|
return (self.n_electrons - self.spin) // 2
|
|
|
|
@property
|
|
def nuclear_repulsion(self) -> float:
|
|
E = 0.0
|
|
for i, a in enumerate(self.atoms):
|
|
for j in range(i + 1, len(self.atoms)):
|
|
b = self.atoms[j]
|
|
R = np.linalg.norm(a.coords - b.coords)
|
|
E += a.Z * b.Z / R
|
|
return E
|
|
|
|
@classmethod
|
|
def from_xyz(cls, xyz: str, charge: int = 0, spin: int = 0, unit: str = 'angstrom'):
|
|
atoms = []
|
|
for line in xyz.strip().splitlines():
|
|
parts = line.split()
|
|
if len(parts) < 4:
|
|
continue
|
|
coords = np.array([float(x) for x in parts[1:4]])
|
|
if unit == 'angstrom':
|
|
coords *= ANGSTROM_TO_BOHR
|
|
atoms.append(Atom(parts[0], coords))
|
|
return cls(atoms, charge, spin)
|
|
|
|
def center(self) -> np.ndarray:
|
|
"""Geometric center of the molecule."""
|
|
return np.mean([a.coords for a in self.atoms], axis=0)
|
|
|
|
def __repr__(self):
|
|
lines = [f"Molecule(charge={self.charge}, spin={self.spin}, n_elec={self.n_electrons})"]
|
|
for a in self.atoms:
|
|
c = a.coords
|
|
lines.append(f" {a.symbol:2s} {c[0]:12.6f} {c[1]:12.6f} {c[2]:12.6f} (Bohr)")
|
|
return "\n".join(lines)
|