Basic Gaussian cube file handler

pycube.py

import os

import numpy as np
from ase import Atoms

class CubeData:
    def __init__(
        self,
        atoms: Atoms,
        origin:np.ndarray,
        n_grids:np.ndarray,
        grid_spacing:np.ndarray,
        data:np.ndarray
    ):
        self.atoms = atoms
        self.origin = origin
        self.n_grids = n_grids
        self.grid_spacing = grid_spacing
        self.data = data

    @classmethod
    def from_file(cls, filename: os.PathLike) -> "CubeData":
        with open(filename, "r") as f:
            # Read comments            
            for _ in range(2):
                next(f)
                
            # line 3: Number of atoms and coordinate of origin
            n_atoms, *origin = map(float, next(f).split())
            n_atoms = int(n_atoms)
            
            # line 4-6: grid spacing
            n_grid_x, *d_grid_x = map(float, next(f).split())
            n_grid_x = int(n_grid_x)
            n_grid_y, *d_grid_y = map(float, next(f).split())
            n_grid_y = int(n_grid_y)
            n_grid_z, *d_grid_z = map(float, next(f).split())
            n_grid_z = int(n_grid_z)
            
            # Assert that the grid spacing is not sheared
            assert (sum(d_grid_x) - d_grid_x[0]) < 1e-6
            assert (sum(d_grid_y) - d_grid_y[1]) < 1e-6
            assert (sum(d_grid_z) - d_grid_z[2]) < 1e-6
            d_grid_x = d_grid_x[0]
            d_grid_y = d_grid_y[1]
            d_grid_z = d_grid_z[2]
            
            # Read data
            atomic_numbers = []
            atom_charges = []
            atom_coords = []
            for _ in range(n_atoms):
                z, q, *R = map(float, next(f).split())
                atomic_numbers.append(int(z))
                atom_charges.append(q)
                atom_coords.append(R)
            data = []
            for line in f:
                numbers = np.array([float(x) for x in line.split()])
                data.append(numbers)
            data = np.concatenate(data).reshape(n_grid_x, n_grid_y, n_grid_z)

        # Construct object
        atoms = Atoms(numbers=atomic_numbers, positions=atom_coords)
        origin = np.asarray(origin)
        n_grids = np.array([n_grid_x, n_grid_y, n_grid_z])
        grid_spacing = np.array([d_grid_x, d_grid_y, d_grid_z])
        return cls(atoms, origin, n_grids, grid_spacing, data)
    
    def write(self, filename: os.PathLike):
        with open(filename, "w") as f:
            # Write comments
            f.write("Generated by CubeData class\n")
            f.write(f"Total {np.prod(self.n_grids)} grid points\n")
            origin_x, origin_y, origin_z = self.origin
            info_line = f"  {len(self.atoms)}  {origin_x:.6f}  {origin_y:.6f}  {origin_z:.6f}\n"
            f.write(info_line)
            # Write grid info
            for i in range(3):
                n = self.n_grids[i]
                spacing = [0, 0, 0]
                spacing[i] = self.grid_spacing[i]
                line = f"  {n}  {spacing[0]:.6f}  {spacing[1]:.6f}  {spacing[2]:.6f}\n"
                f.write(line)
            for i in range(len(self.atoms)):
                Z = self.atoms[i].number
                x, y, z = self.atoms[i].position
                line = f"  {Z}  {float(Z):.6f}  {x:.6f}  {y:.6f}  {z:.6f}\n"
                f.write(line)
            for i in range(self.n_grids[0]):
                for j in range(self.n_grids[1]):
                    for k in range(self.n_grids[2]):
                        f.write(f"{self.data[i, j, k]:.6E}")
                        if k % 6 == 5:
                            f.write("\n")
                        else:
                            f.write("  ")
                    f.write("\n")
    
    def __neg__(self) -> "CubeData":
        new_data = -self.data
        return CubeData(self.atoms.copy(), self.origin.copy(), self.n_grids.copy(), self.grid_spacing.copy(), new_data)
    
    def __add__(self, other: "CubeData") -> "CubeData":
        new_data = self.data + other.data
        return CubeData(self.atoms.copy(), self.origin.copy(), self.n_grids.copy(), self.grid_spacing.copy(), new_data)
    
    def __sub__(self, other: "CubeData") -> "CubeData":
        new_data = self.data - other.data
        return CubeData(self.atoms.copy(), self.origin.copy(), self.n_grids.copy(), self.grid_spacing.copy(), new_data)
    
    def __iadd__(self, other: "CubeData") -> "CubeData":
        self.data += other.data
        return self

    def __isub__(self, other: "CubeData") -> "CubeData":
        self.data -= other.data
        return self