"""Structure storage for the eOn Python server.
Canonical on-disk form is ``readcon.ConFrame`` (``.con`` via the readcon package).
In-memory working form is :class:`Structure`: contiguous numpy arrays for geometry
algorithms (PBC, neighbor lists, displacements). This replaces the historical
aselite-era mini-Atoms type without requiring ASE.
``Atoms`` is an alias of :class:`Structure` for call-site compatibility.
"""
from __future__ import annotations
from typing import List, Optional, Sequence, Union
import numpy as np
import readcon
from eon.geometry.cell import box_to_length_angle, length_angle_to_box
[docs]
class Structure:
"""Mutable atomic configuration (numpy-backed).
Attributes
----------
r : (N, 3) float
Cartesian positions.
free : (N,) float
1.0 if free to move, 0.0 if fixed (eOn convention).
box : (3, 3) float
Cell matrix with lattice vectors as **rows** (ASE/vesin convention).
names : list[str]
Chemical symbols, length N.
mass : (N,) float
Atomic masses.
"""
__slots__ = ("r", "free", "box", "names", "mass")
def __init__(self, n_atoms: int = 0):
self.r = np.zeros((n_atoms, 3), dtype=float)
self.free = np.ones(n_atoms, dtype=float)
self.box = np.zeros((3, 3), dtype=float)
self.names: List[str] = [""] * n_atoms
self.mass = np.zeros(n_atoms, dtype=float)
def __len__(self) -> int:
return int(self.r.shape[0])
[docs]
def copy(self) -> "Structure":
p = Structure(len(self))
p.r = self.r.copy()
p.free = self.free.copy()
p.box = self.box.copy()
p.names = list(self.names)
p.mass = self.mass.copy()
return p
[docs]
def free_r(self) -> np.ndarray:
"""Positions of free (unfixed) atoms only."""
return self.r[np.asarray(self.free, dtype=bool)]
[docs]
def free_mask(self) -> np.ndarray:
return np.asarray(self.free, dtype=bool)
[docs]
def fixed_mask(self) -> np.ndarray:
return ~self.free_mask()
[docs]
def append(self, r, free, name, mass) -> None:
self.r = np.append(self.r, [r], 0)
self.free = np.append(self.free, free)
self.names.append(name)
self.mass = np.append(self.mass, mass)
# --- readcon ConFrame bridge (canonical I/O type) ---
[docs]
@classmethod
def from_conframe(cls, frame: "readcon.ConFrame") -> "Structure":
"""Build a Structure from a readcon ConFrame (live API)."""
n = len(frame)
p = cls(n)
boxlengths = np.asarray(list(frame.cell), dtype=float)
boxangles = np.asarray(list(frame.angles), dtype=float)
p.box = length_angle_to_box(boxlengths, boxangles)
# Prefer bulk coords when available
try:
coords = np.asarray(frame.coords_array(), dtype=float)
if coords.shape == (n, 3):
p.r = coords.copy()
else:
raise ValueError("coords shape")
except Exception:
for i, atom in enumerate(frame.atoms):
p.r[i] = [atom.x, atom.y, atom.z]
for i, atom in enumerate(frame.atoms):
p.names[i] = atom.symbol
p.mass[i] = atom.mass if atom.mass is not None else 0.0
fixed = atom.fixed
p.free[i] = 0.0 if (fixed is not None and any(fixed)) else 1.0
return p
[docs]
def to_conframe(
self,
prebox_header: Optional[Sequence[str]] = None,
postbox_header: Optional[Sequence[str]] = None,
) -> "readcon.ConFrame":
"""Convert to a readcon ConFrame for writing."""
lengths, angles = box_to_length_angle(self.box)
atom_list = []
for i in range(len(self)):
is_fixed = bool(self.free[i] == 0)
atom_list.append(
readcon.Atom(
symbol=self.names[i],
x=float(self.r[i][0]),
y=float(self.r[i][1]),
z=float(self.r[i][2]),
fixed=[is_fixed, is_fixed, is_fixed],
atom_id=i + 1,
mass=float(self.mass[i]),
)
)
return readcon.ConFrame(
cell=list(lengths),
angles=list(angles),
atoms=atom_list,
prebox_header=list(prebox_header)
if prebox_header is not None
else ["Generated by eOn", ""],
postbox_header=list(postbox_header) if postbox_header is not None else None,
)
# Back-compat name used throughout the server
Atoms = Structure
[docs]
def structure_to_ase(structure: "Structure", *, pbc: bool = True):
"""Optional ASE export (requires ase + pyeonclient.ase_bridge)."""
from pyeonclient.ase_bridge import structure_to_ase as _f
return _f(structure, pbc=pbc)
[docs]
def ase_to_structure(atoms):
"""Optional ASE import (requires ase + pyeonclient.ase_bridge)."""
from pyeonclient.ase_bridge import ase_to_structure as _f
return _f(atoms)
# Methods on Structure for ergonomics
def _structure_to_ase(self, *, pbc: bool = True):
return structure_to_ase(self, pbc=pbc)
Structure.to_ase = _structure_to_ase # type: ignore[attr-defined]
def _structure_from_ase(cls, atoms):
return ase_to_structure(atoms)
Structure.from_ase = classmethod(_structure_from_ase) # type: ignore[attr-defined, assignment]