crisp-ase 1.0.0.post0.dev0__py3-none-any.whl

CRISP/simulation_utility/interatomic_distances.py

@@ -0,0 +1,198 @@
+"""
+CRISP/simulation_utility/interatomic_distances.py
+
+This module provides tools to calculate and analyze interatomic distances from
+molecular dynamics trajectories.
+"""
+
+import os
+import numpy as np
+import pickle
+from ase.io import read
+from ase import Atoms
+from typing import Union, Tuple, List, Dict, Any
+from joblib import Parallel, delayed
+
+
+def indices(atoms: Atoms, ind: Union[str, List[Union[int, str]], None]) -> np.ndarray:
+    """Extract atom indices from various input types.
+    
+    Parameters
+    ----------
+    atoms : ase.Atoms
+        Atoms object containing atomic coordinates and elements
+    ind : str, list, or None
+        Specification for which atoms to select:
+        - "all" or None: all atoms
+        - string ending with ".npy": load indices from NumPy file
+        - list of integers: direct atom indices
+        - list of strings: chemical symbols to select
+        
+    Returns
+    -------
+    np.ndarray
+        Array of atom indices
+        
+    Raises
+    ------
+    ValueError
+        If the index type is not recognized
+    """
+    if ind == "all" or ind is None:
+        return np.arange(len(atoms))
+    if isinstance(ind, str) and ind.endswith(".npy"):
+        return np.load(ind, allow_pickle=True)
+    if not isinstance(ind, list):
+        ind = [ind]
+    if any(isinstance(item, int) for item in ind):
+        return np.array(ind)
+    if any(isinstance(item, str) for item in ind):
+        idx = []
+        for symbol in ind:
+            idx.append(np.where(np.array(atoms.get_chemical_symbols()) == symbol)[0])
+        return np.concatenate(idx)
+    raise ValueError("Invalid index type")
+
+
+def distance_calculation(
+    traj_path: str,
+    frame_skip: int,
+    index_type: Union[str, List[Union[int, str]]] = "all"
+) -> Tuple[List[np.ndarray], List[np.ndarray]]:
+    """Calculate distance matrices for multiple frames in a trajectory.
+    
+    Parameters
+    ----------
+    traj_path : str
+        Path to the trajectory file in any format supported by ASE
+    frame_skip : int
+        Read every nth frame (n=frame_skip)
+    index_type : str, list, or None, optional
+        Specification for which atoms to select for sub-matrix (default: "all")
+        
+    Returns
+    -------
+    Tuple[List[np.ndarray], List[np.ndarray]]
+        Two lists containing:
+        1. Full distance matrices for all frames
+        2. Sub-matrices for specified atoms
+        
+    Raises
+    ------
+    ValueError
+        If no frames were found in the trajectory or if format is unsupported
+    """
+    try:
+        # Let ASE auto-detect file format based on extension
+        frames = read(traj_path, index=f"::{frame_skip}")
+        
+        # Handle the case when a single frame is returned (not a list)
+        if not isinstance(frames, list):
+            frames = [frames]
+            
+        if not frames:
+            raise ValueError("No frames were found in the trajectory using the given frame_skip.")
+
+        def process_frame(frame: Atoms) -> Tuple[np.ndarray, np.ndarray]:
+            dm = frame.get_all_distances(mic=True)
+            idx = indices(frame, index_type)
+            sub_dm = dm[np.ix_(idx, idx)]
+            return dm, sub_dm
+
+        results = Parallel(n_jobs=-1)(delayed(process_frame)(frame) for frame in frames)
+        full_dms, sub_dms = zip(*results)
+        return list(full_dms), list(sub_dms)
+        
+    except ValueError as e:
+        raise e
+    except Exception as e:
+        raise ValueError(f"Error processing trajectory: {e}. Check if the format is supported by ASE.")
+
+
+def save_distance_matrices(
+    full_dms: List[np.ndarray], 
+    sub_dms: List[np.ndarray],
+    index_type: Union[str, List[Union[int, str]]] = "all",
+    output_dir: str = "distance_calculations"
+) -> None:
+    """Save distance matrices to pickle file.
+    
+    Parameters
+    ----------
+    full_dms : List[np.ndarray]
+        List of full distance matrices
+    sub_dms : List[np.ndarray]
+        List of sub-matrices for specified atoms
+    index_type : str, list, or None, optional
+        Type of index selection used (default: "all")
+    output_dir : str, optional
+        Directory to save output file (default: "distance_calculations")
+        
+    Returns
+    -------
+    None
+        Saves results to disk
+    """
+    data = {"full_dms": full_dms}
+    if index_type not in ["all", None]:
+        data["sub_dms"] = sub_dms
+        
+    os.makedirs(output_dir, exist_ok=True)
+    output_path = os.path.join(output_dir, "distance_matrices.pkl")
+    with open(output_path, "wb") as f:
+        pickle.dump(data, f)
+    print(f"Distance matrices saved in '{output_path}'")
+
+
+def calculate_interatomic_distances(
+    traj_path: str,
+    frame_skip: int = 10,
+    index_type: Union[str, List[Union[int, str]]] = "all",
+    output_dir: str = "distance_calculations",
+    save_results: bool = True
+) -> Dict[str, List[np.ndarray]]:
+    """
+    Calculate interatomic distances for a trajectory and optionally save results.
+    
+    Parameters
+    ----------
+    traj_path : str
+        Path to the trajectory file
+    frame_skip : int, optional
+        Read every nth frame (default: 10)
+    index_type : str, list, or None, optional
+        Specification for which atoms to select (default: "all")
+    output_dir : str, optional
+        Directory to save output file (default: "distance_calculations")
+    save_results : bool, optional
+        Whether to save results to disk (default: True)
+        
+    Returns
+    -------
+    Dict[str, List[np.ndarray]]
+        Dictionary containing full distance matrices and optionally sub-matrices
+        
+    Examples
+    --------
+    >>> results = calculate_interatomic_distances("trajectory.traj")
+    >>> first_frame_distances = results["full_dms"][0]
+    >>> print(f"Distance matrix shape: {first_frame_distances.shape}")
+    """
+    print(f"Calculating interatomic distances from '{traj_path}'")
+    print(f"Using frame skip: {frame_skip}")
+    print(f"Index type: {index_type}")
+    
+    full_dms, sub_dms = distance_calculation(traj_path, frame_skip, index_type)
+    
+    print(f"Processed {len(full_dms)} frames")
+    print(f"Full matrix shape: {full_dms[0].shape}")
+    print(f"Sub-matrix shape: {sub_dms[0].shape}")
+    
+    results = {"full_dms": full_dms}
+    if index_type not in ["all", None]:
+        results["sub_dms"] = sub_dms
+    
+    if save_results:
+        save_distance_matrices(full_dms, sub_dms, index_type, output_dir)
+    
+    return results

CRISP/simulation_utility/subsampling.py

@@ -0,0 +1,221 @@
+"""
+CRISP/simulation_utility/subsampling.py
+
+This module provides functionality for structure subsampling from molecular dynamics
+trajectories using Farthest Point Sampling (FPS) with SOAP descriptors.
+"""
+
+import numpy as np
+from ase.io import read, write
+import fpsample
+import glob
+import os
+from dscribe.descriptors import SOAP
+import matplotlib.pyplot as plt
+from joblib import Parallel, delayed
+from typing import Union, List
+
+
+def indices(atoms, ind: Union[str, List[Union[int, str]]]) -> np.ndarray:
+    """
+    Extract atom indices from an ASE Atoms object based on the input specifier.
+    
+    Parameters
+    ----------
+    atoms : ase.Atoms
+        ASE Atoms object containing atomic structure
+    ind : Union[str, List[Union[int, str]]]
+        Index specifier, can be:
+        - "all" or None: all atoms
+        - string ending with ".npy": load indices from NumPy file
+        - integer or list of integers: direct atom indices
+        - string or list of strings: chemical symbols to select
+        
+    Returns
+    -------
+    np.ndarray
+        Array of selected indices
+        
+    Raises
+    ------
+    ValueError
+        If the index type is invalid
+    """
+    # Select all atoms
+    if ind == "all" or ind is None:
+        return np.arange(len(atoms))
+    
+    # Load from NumPy file
+    if isinstance(ind, str) and ind.endswith(".npy"):
+        return np.load(ind, allow_pickle=True)
+    
+    # Convert single items to list
+    if not isinstance(ind, list):
+        ind = [ind]
+    
+    # Handle integer indices directly
+    if any(isinstance(item, int) for item in ind):
+        return np.array(ind)
+    
+    # Handle chemical symbols
+    if any(isinstance(item, str) for item in ind):
+        idx = []
+        for symbol in ind:
+            idx.append(np.where(np.array(atoms.get_chemical_symbols()) == symbol)[0])
+        return np.concatenate(idx)
+    
+    raise ValueError("Invalid index type")
+
+
+def compute_soap(structure, all_spec, rcut, idx):
+    """Compute SOAP descriptors for a given structure.
+    
+    Parameters
+    ----------
+    structure : ase.Atoms
+        Atomic structure for which to compute SOAP descriptors
+    all_spec : list
+        List of chemical elements to include in the descriptor
+    rcut : float
+        Cutoff radius for the SOAP descriptor in Angstroms
+    idx : numpy.ndarray
+        Indices of atoms to use as centers for SOAP calculation
+        
+    Returns
+    -------
+    numpy.ndarray
+        Average SOAP descriptor vector for the structure
+    """
+    periodic_cell = structure.cell.volume > 0
+    soap = SOAP(
+        species=all_spec,
+        periodic=periodic_cell,
+        r_cut=rcut,
+        n_max=8,
+        l_max=6,
+        sigma=0.5,
+        sparse=False
+    )
+    soap_ind = soap.create(structure, centers=idx)
+    return np.mean(soap_ind, axis=0)
+
+
+def create_repres(traj_path, rcut=6, ind="all", n_jobs=-1):
+    """Create SOAP representation vectors for a trajectory.
+    
+    Parameters
+    ----------
+    traj_path : list
+        List of ase.Atoms objects representing a trajectory
+    rcut : float, optional
+        Cutoff radius for the SOAP descriptor in Angstroms (default: 6)
+    ind : str, list, or None, optional
+        Specification for which atoms to use as SOAP centers (default: "all")
+    n_jobs : int, optional
+        Number of parallel jobs to run; -1 uses all available cores (default: -1)
+        
+    Returns
+    -------
+    numpy.ndarray
+        Array of SOAP descriptors for each frame in the trajectory
+    """
+    all_spec = traj_path[0].get_chemical_symbols()
+    idx = indices(traj_path[0], ind=ind)
+
+    repres = Parallel(n_jobs=n_jobs)(
+        delayed(compute_soap)(structure, all_spec, rcut, idx) for structure in traj_path
+    )
+
+    return np.array(repres)
+
+
+def subsample(traj_path, n_samples=50, index_type="all", rcut=6.0, file_format=None, 
+             plot_subsample=False, frame_skip=1, output_dir="subsampled_structures"):
+    """Subsample a trajectory using Farthest Point Sampling with SOAP descriptors.
+    
+    Parameters
+    ----------
+    traj_path : str
+        Path pattern to trajectory file(s); supports globbing
+    n_samples : int, optional
+        Number of frames to select (default: 50)
+    index_type : str, list, or None, optional
+        Specification for which atoms to use for SOAP calculation (default: "all")
+    rcut : float, optional
+        Cutoff radius for SOAP in Angstroms (default: 6.0)
+    file_format : str, optional
+        File format for ASE I/O (default: None, auto-detect)
+    plot_subsample : bool, optional
+        Whether to generate a plot of FPS distances (default: False)
+    frame_skip : int, optional
+        Read every nth frame from the trajectory (default: 1)
+    output_dir : str, optional
+        Directory to save the subsampled structures (default: "subsampled_structures")
+        
+    Returns
+    -------
+    list
+        List of selected ase.Atoms frames
+        
+    Notes
+    -----
+    The selected frames and plots are saved in the specified output directory
+    """
+    traj_files = glob.glob(traj_path)
+    
+    # Check if any matching files were found
+    if not traj_files:
+        raise ValueError(f"No files found matching pattern: {traj_path}")
+
+    trajec = []
+    for file in traj_files:
+        if file_format is not None:
+            trajec += read(file, index=f'::{frame_skip}', format=file_format)
+        else:
+            trajec += read(file, index=f'::{frame_skip}')
+
+    if not isinstance(trajec, list): 
+        trajec = [trajec]
+    
+    repres = create_repres(trajec, ind=index_type, rcut=rcut)
+    
+    # Ensure we don't request more samples than available frames
+    n_samples = min(n_samples, len(trajec))
+    
+    perm = fpsample.fps_sampling(repres, n_samples, start_idx=0)
+
+    fps_frames = []
+
+    for str_idx, frame in enumerate(perm):
+        new_frame = trajec[frame]
+        fps_frames.append(new_frame)
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    if plot_subsample:
+        distance = []
+        for i in range(1, len(perm)):
+            distance.append(np.min(np.linalg.norm(repres[perm[:i]] - repres[perm[i]], axis=1)))
+
+        plt.figure(figsize=(8, 6))
+        plt.plot(distance, c="blue", linewidth=2)
+        plt.ylim([0, 1.1 * max(distance)])
+        plt.xlabel("Number of subsampled structures")
+        plt.ylabel("Euclidean distance")
+        plt.title("FPS Subsampling")
+        plt.savefig(os.path.join(output_dir, "subsampled_convergence.png"), dpi=300)
+        plt.show()
+        plt.close()
+        print(f"Saved convergence plot to {os.path.join(output_dir, 'subsampled_convergence.png')}")
+
+    # Extract the base filename without path for output file using os.path for platform independence
+    base_filename = os.path.basename(traj_files[0])
+    output_file = os.path.join(output_dir, f"subsample_{base_filename}")
+    
+    try:
+        write(output_file, fps_frames, format=file_format)
+        print(f"Saved {len(fps_frames)} subsampled structures to {output_file}")
+    except Exception as e:
+        print(f"Error saving subsampled structures: {e}")
+
+    return fps_frames

CRISP/tests/__init__.py

@@ -0,0 +1,3 @@
+"""
+Empty init file in case you choose a package besides PyTest such as Nose which may look for such a file.
+"""

CRISP/tests/test_CRISP.py

@@ -0,0 +1,28 @@
+"""Basic CRISP package tests."""
+
+import pytest
+import CRISP
+
+
+def test_crisp_import():
+    """Test that CRISP can be imported."""
+    assert CRISP is not None
+
+
+def test_crisp_version():
+    """Test that CRISP has a version."""
+    try:
+        version = CRISP.__version__
+        assert isinstance(version, str)
+    except AttributeError:
+        pass
+
+
+def test_crisp_modules():
+    """Test that main modules can be imported."""
+    try:
+        from CRISP.simulation_utility import atomic_indices
+        from CRISP.data_analysis import msd
+        assert True
+    except ImportError as e:
+        pytest.fail(f"Failed to import CRISP modules: {e}")