matlanticgrrm.io.formats のソースコード

from ase.io import iread, read, write
from ase import Atoms
import pickle
from pathlib import Path

def lread(traj,pkl=None,status=[0]):
    if pkl:
        with open(pkl,"rb") as f:
            l,n = pickle.load(f)
        if n == 1:
            return _read_with_status(traj,l,status)
        else:
            return _read_images(traj,l,n,status)
    else:
        """trajファイルのみでatomsリストを再生する.statusを指定できない."""
        atoms_list = _read_without_status(traj)
        if len(atoms_list) == 1 :
            return atoms_list[0]
        else:
            return atoms_list

def _read_without_status(traj):
    return [atoms if atoms else None for atoms in iread(traj)]

def _read_with_status(traj,l,status):
    return [atoms if i in status else None for i, atoms in zip(l,iread(traj))]

def _read_images(traj,l,n,status):
    gen = iread(traj) # atomsを返すジェネレーター
    imagess = []
    for i in l:
        images = []
        for _ in range(n):
            if i in status: # 収束済み以外はNoneにする
                images.append(next(gen))
            else:
                next(gen)
                images = None
        imagess.append(images)
    return imagess
        
def lwrite(filename,atoms_list):
    """atoms_listからtrajファイルを作成する
    
    通常はASEのwriteを使用すれば良いが, listの要素にNoneが含まれている場合この関数を使用する.
    Noneの要素の部分には空のAtomsを挿入される.
    
    Parameters:
    
    filename: str or path object
        trajファイルのパス
    atoms_list: str or path object
        保存するAomtsリストの1次元リストまたは2次元リスト
    
    Note:
        Noneの要素の部分には空のAtomsを挿入される. 空のAtomsはASEの仕様上, 真偽値でFalseを返すので
        lwriteで作成したtrajファイルからatoms_listを再生するには
        
        >>> from ase.io import iread
        >>> atoms_list =z [atoms if atoms else None for atoms in iread("*traj")]
        
        とすれば良い
        
    Note:
        atoms_listがimages(Atomsの2次元リスト)の場合はfilenameと同じ名前のpickleファイルも作成される.
    
    """
    for atoms in atoms_list:
        if atoms is not None:
            if type(atoms) == list:
                n_images = len(atoms)
                _write_images_list(filename,atoms_list,n_images)
            else:
                _write_atoms_list(filename,atoms_list)
        break
        
def _write_atoms_list(filename,atoms_list):
    atoms_list = [atoms if atoms else Atoms() for atoms in atoms_list]
    write(filename,atoms_list)

def _write_images_list(filename,atoms_list,n_images):
    flat_atoms_list = []
    status_list = []
    for images in atoms_list:
        if images:
            for atoms in images:
                flat_atoms_list.append(atoms)
            status_list.append(0)
        else:
            for _ in range(n_images):
                flat_atoms_list.append(Atoms())
            status_list.append(2)
    write(filename,flat_atoms_list)
    
    with open(f"{Path(filename).name}.pickle","wb") as f:
        pickle.dump((status_list,n_images),f)

def traj2cif(traj):
    traj_path = Path(traj)
    atoms = read(traj_path)
    write(traj_path.with_suffix('.cif'),atoms)
    
def traj2xyz(traj):
    traj_path = Path(traj)
    atoms = read(traj_path)
    write(traj_path.with_suffix('.xyz'),atoms)