crisp-ase 1.1.2__py3-none-any.whl

CRISP/tests/DataAnalysis/test_msd_extended.py

@@ -0,0 +1,522 @@
+"""Extended comprehensive tests for Mean Square Displacement (MSD) module."""
+import pytest
+import numpy as np
+import pandas as pd
+import os
+import tempfile
+import shutil
+from ase import Atoms
+from ase.io import write, read
+from CRISP.data_analysis.msd import (
+    calculate_msd,
+    calculate_frame_msd,
+    read_trajectory_chunk,
+    save_msd_data,
+    plot_diffusion_time_series,
+    block_averaging_error,
+    calculate_msd_windowed,
+    calculate_frame_msd_windowed,
+    msd_analysis,
+    calculate_diffusion_coefficient,
+)
+
+
+class TestCalculateFrameMSD:
+    """Test single frame MSD calculation."""
+    
+    def test_calculate_frame_msd_basic(self):
+        """Test basic frame MSD calculation."""
+        ref_frame = Atoms('H2O', positions=[[0, 0, 0], [1, 0, 0], [0, 1, 0]])
+        curr_frame = Atoms('H2O', positions=[[0.1, 0, 0], [1.1, 0, 0], [0.1, 1, 0]])
+        atom_indices = np.array([0, 1, 2])
+        
+        frame_idx, msd = calculate_frame_msd(0, curr_frame, ref_frame, atom_indices, msd_direction=False)
+        
+        assert frame_idx == 0
+        assert isinstance(msd, (float, np.floating))
+        assert msd > 0
+    
+    def test_calculate_frame_msd_directional(self):
+        """Test directional MSD calculation."""
+        ref_frame = Atoms('H2O', positions=[[0, 0, 0], [1, 0, 0], [0, 1, 0]])
+        curr_frame = Atoms('H2O', positions=[[0.1, 0.2, 0.3], [1.1, 0, 0], [0.1, 1, 0]])
+        atom_indices = np.array([0, 1, 2])
+        
+        result = calculate_frame_msd(1, curr_frame, ref_frame, atom_indices, msd_direction=True)
+        
+        assert len(result) == 4  # frame_idx, msd_x, msd_y, msd_z
+        frame_idx, msd_x, msd_y, msd_z = result
+        assert frame_idx == 1
+        assert isinstance(msd_x, (float, np.floating))
+        assert isinstance(msd_y, (float, np.floating))
+        assert isinstance(msd_z, (float, np.floating))
+        assert all(v >= 0 for v in [msd_x, msd_y, msd_z])
+    
+    def test_calculate_frame_msd_single_atom(self):
+        """Test MSD for single atom."""
+        ref_frame = Atoms('H', positions=[[0, 0, 0]])
+        curr_frame = Atoms('H', positions=[[1, 1, 1]])
+        atom_indices = np.array([0])
+        
+        frame_idx, msd = calculate_frame_msd(0, curr_frame, ref_frame, atom_indices)
+        
+        assert msd == pytest.approx(3.0)  # (1^2 + 1^2 + 1^2) / 1
+    
+    def test_calculate_frame_msd_no_displacement(self):
+        """Test MSD when atoms don't move."""
+        frame = Atoms('H2O', positions=[[0, 0, 0], [1, 0, 0], [0, 1, 0]])
+        atom_indices = np.array([0, 1, 2])
+        
+        frame_idx, msd = calculate_frame_msd(0, frame, frame, atom_indices)
+        
+        assert msd == pytest.approx(0.0)
+
+
+class TestReadTrajectoryChunk:
+    """Test trajectory file reading."""
+    
+    @pytest.fixture
+    def sample_trajectory(self):
+        """Create a sample trajectory file."""
+        temp_dir = tempfile.mkdtemp()
+        traj_path = os.path.join(temp_dir, 'test_traj.traj')
+        
+        # Create multi-frame trajectory
+        frames = []
+        for i in range(5):
+            atoms = Atoms('H2O', positions=[[0, 0, 0], [1, 0, 0], [0, 1, 0]])
+            atoms.positions += np.random.rand(3, 3) * 0.1 * i
+            frames.append(atoms)
+        
+        write(traj_path, frames)
+        yield traj_path
+        shutil.rmtree(temp_dir)
+    
+    def test_read_trajectory_chunk_all_frames(self, sample_trajectory):
+        """Test reading all frames."""
+        frames = read_trajectory_chunk(sample_trajectory, ':')
+        
+        assert isinstance(frames, list)
+        assert len(frames) == 5
+        assert all(isinstance(f, Atoms) for f in frames)
+    
+    def test_read_trajectory_chunk_specific_slice(self, sample_trajectory):
+        """Test reading specific frame slice."""
+        frames = read_trajectory_chunk(sample_trajectory, '0:3')
+        
+        assert len(frames) == 3
+    
+    def test_read_trajectory_chunk_with_frame_skip(self, sample_trajectory):
+        """Test reading with frame skip."""
+        frames = read_trajectory_chunk(sample_trajectory, ':', frame_skip=2)
+        
+        # 5 frames with skip=2 should give 3 frames (indices 0, 2, 4)
+        assert len(frames) == 3
+    
+    def test_read_trajectory_chunk_invalid_file(self):
+        """Test reading non-existent file."""
+        result = read_trajectory_chunk('/nonexistent/path.traj', ':')
+        
+        assert result == []
+
+
+class TestCalculateMSD:
+    """Test MSD calculation for trajectories."""
+    
+    @pytest.fixture
+    def trajectory(self):
+        """Create sample trajectory."""
+        frames = []
+        for i in range(10):
+            atoms = Atoms('H2O', positions=[[0, 0, 0], [1, 0, 0], [0, 1, 0]])
+            # Linear displacement
+            atoms.positions += np.array([0.01*i, 0.01*i, 0])
+            frames.append(atoms)
+        return frames
+    
+    def test_calculate_msd_basic(self, trajectory):
+        """Test basic MSD calculation."""
+        timestep = 1.0
+        atom_indices = np.array([0, 1, 2])
+        
+        result = calculate_msd(trajectory, timestep, atom_indices=atom_indices)
+        
+        assert isinstance(result, tuple)
+        assert len(result) == 2  # (msd_values, msd_times)
+        msd_values, msd_times = result
+        assert len(msd_values) > 0
+        assert len(msd_times) > 0
+        assert msd_times[0] == 0.0
+    
+    def test_calculate_msd_directional(self, trajectory):
+        """Test directional MSD calculation."""
+        timestep = 1.0
+        atom_indices = np.array([0, 1, 2])
+        
+        result = calculate_msd(trajectory, timestep, atom_indices=atom_indices, msd_direction=True)
+        
+        assert isinstance(result, tuple)
+        assert len(result) == 4  # (msd_times, msd_x, msd_y, msd_z)
+        msd_times, msd_x, msd_y, msd_z = result
+        assert all(len(arr) > 0 for arr in [msd_times, msd_x, msd_y, msd_z])
+    
+    def test_calculate_msd_with_ignore_frames(self, trajectory):
+        """Test MSD calculation ignoring initial frames."""
+        timestep = 1.0
+        atom_indices = np.array([0, 1, 2])
+        
+        result_full = calculate_msd(trajectory, timestep, atom_indices=atom_indices, ignore_n_images=0)
+        result_skip = calculate_msd(trajectory, timestep, atom_indices=atom_indices, ignore_n_images=3)
+        
+        msd_full, time_full = result_full
+        msd_skip, time_skip = result_skip
+        
+        assert len(time_skip) < len(time_full)
+    
+    def test_calculate_msd_single_atom_index(self, trajectory):
+        """Test MSD for single atom."""
+        timestep = 1.0
+        atom_indices = np.array([0])
+        
+        msd_values, msd_times = calculate_msd(trajectory, timestep, atom_indices=atom_indices)
+        
+        assert isinstance(msd_values, np.ndarray)
+        assert len(msd_values) > 0
+        # MSD should increase with time for linearly moving particle
+        assert msd_values[-1] > msd_values[0]
+
+
+class TestSaveMSDData:
+    """Test MSD data saving."""
+    
+    def test_save_msd_data_basic(self):
+        """Test basic MSD data saving."""
+        temp_dir = tempfile.mkdtemp()
+        
+        # Use tuple format: (msd_values, msd_times)
+        msd_values = np.array([0, 0.1, 0.2, 0.3, 0.4])
+        msd_times = np.array([0, 1, 2, 3, 4])
+        msd_data = (msd_values, msd_times)
+        
+        csv_file = os.path.join(temp_dir, 'msd_test.csv')
+        save_msd_data(msd_data, csv_file, output_dir=temp_dir)
+        
+        assert os.path.exists(csv_file)
+        
+        # Read back and verify
+        df = pd.read_csv(csv_file)
+        assert len(df) > 0
+        
+        shutil.rmtree(temp_dir)
+    
+    def test_save_msd_data_creates_directory(self):
+        """Test that output directory is created."""
+        temp_dir = tempfile.mkdtemp()
+        output_subdir = os.path.join(temp_dir, 'new_output')
+        
+        # Use tuple format: (msd_values, msd_times)
+        msd_values = np.array([0, 0.1, 0.2])
+        msd_times = np.array([0, 1, 2])
+        msd_data = (msd_values, msd_times)
+        
+        csv_file = os.path.join(output_subdir, 'msd.csv')
+        save_msd_data(msd_data, csv_file, output_dir=output_subdir)
+        
+        assert os.path.exists(output_subdir)
+        
+        shutil.rmtree(temp_dir)
+
+
+class TestPlotDiffusionTimeSeries:
+    """Test diffusion time series plotting."""
+    
+    def test_plot_diffusion_time_series_basic(self):
+        """Test basic plotting functionality."""
+        temp_dir = tempfile.mkdtemp()
+        
+        msd_times = np.array([0, 1, 2, 3, 4, 5])
+        msd_values = np.array([0, 0.1, 0.2, 0.3, 0.4, 0.5])
+        
+        # Plot without saving (just test it doesn't crash)
+        plot_diffusion_time_series(msd_times, msd_values, min_window=2)
+        
+        shutil.rmtree(temp_dir)
+    
+    def test_plot_diffusion_with_intercept(self):
+        """Test plotting with intercept."""
+        msd_times = np.array([0, 1, 2, 3, 4, 5])
+        msd_values = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6])
+        
+        # Should not raise error
+        plot_diffusion_time_series(msd_times, msd_values, with_intercept=True, min_window=2)
+    
+    def test_plot_diffusion_dimension_variations(self):
+        """Test plotting with different dimensions."""
+        msd_times = np.array([0, 1, 2, 3, 4])
+        msd_values = np.array([0, 0.1, 0.2, 0.3, 0.4])
+        
+        for dim in [1, 2, 3]:
+            plot_diffusion_time_series(msd_times, msd_values, dimension=dim, min_window=2)
+
+
+class TestBlockAveragingError:
+    """Test block averaging error analysis."""
+    
+    def test_block_averaging_error_basic(self):
+        """Test basic block averaging."""
+        msd_times = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
+        msd_values = np.array([0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9])
+        
+        result = block_averaging_error(msd_times, msd_values, n_blocks=5)
+        
+        assert result is not None
+        # Should return tuple or similar structure
+        assert isinstance(result, (tuple, dict, type(None)))
+    
+    def test_block_averaging_different_block_counts(self):
+        """Test with different block counts."""
+        msd_times = np.arange(0, 20, 1.0)
+        msd_values = msd_times * 0.1
+        
+        for n_blocks in [2, 3, 5, 10]:
+            result = block_averaging_error(msd_times, msd_values, n_blocks=n_blocks)
+            assert result is not None
+    
+    def test_block_averaging_dimension_variations(self):
+        """Test block averaging with different dimensions."""
+        msd_times = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
+        msd_values = np.array([0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9])
+        
+        for dim in [1, 2, 3]:
+            result = block_averaging_error(msd_times, msd_values, dimension=dim, n_blocks=2)
+            assert result is not None
+
+
+class TestCalculateFrameMSDWindowed:
+    """Test windowed frame MSD calculation."""
+    
+    def test_calculate_frame_msd_windowed_basic(self):
+        """Test basic windowed MSD calculation."""
+        pos_i = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]])
+        pos_j = np.array([[0.1, 0, 0], [1.1, 0, 0], [0.1, 1, 0]])
+        atom_indices = np.array([0, 1, 2])
+        
+        result = calculate_frame_msd_windowed(pos_i, pos_j, atom_indices, msd_direction=False)
+        
+        assert isinstance(result, (float, np.floating))
+        assert result > 0
+    
+    def test_calculate_frame_msd_windowed_directional(self):
+        """Test windowed directional MSD."""
+        pos_i = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]])
+        pos_j = np.array([[0.1, 0.2, 0.3], [1.1, 0, 0], [0.1, 1, 0]])
+        atom_indices = np.array([0, 1, 2])
+        
+        result = calculate_frame_msd_windowed(pos_i, pos_j, atom_indices, msd_direction=True)
+        
+        assert isinstance(result, tuple)
+        assert len(result) == 3  # msd_x, msd_y, msd_z
+
+
+class TestCalculateMSDWindowed:
+    """Test windowed MSD calculation for trajectories."""
+    
+    @pytest.fixture
+    def window_trajectory(self):
+        """Create trajectory for windowed tests."""
+        frames = []
+        for i in range(15):
+            atoms = Atoms('H2O', positions=[[0, 0, 0], [1, 0, 0], [0, 1, 0]])
+            atoms.positions += np.array([0.05*i, 0.05*i, 0])
+            frames.append(atoms)
+        return frames
+    
+    def test_calculate_msd_windowed_basic(self, window_trajectory):
+        """Test basic windowed MSD calculation."""
+        timestep = 1.0
+        atom_indices = np.array([0, 1, 2])
+        
+        # Test without window_size parameter (not supported)
+        result = calculate_msd_windowed(
+            window_trajectory, 
+            timestep, 
+            atom_indices=atom_indices
+        )
+        
+        assert result is not None
+        assert len(result) == 2
+
+
+class TestMSDAnalysis:
+    """Test complete MSD analysis workflow."""
+    
+    @pytest.fixture
+    def test_trajectory_file(self):
+        """Create test trajectory file."""
+        temp_dir = tempfile.mkdtemp()
+        traj_path = os.path.join(temp_dir, 'test_analysis.traj')
+        
+        frames = []
+        for i in range(10):
+            atoms = Atoms('H2O', positions=[[0, 0, 0], [1, 0, 0], [0, 1, 0]])
+            atoms.positions += np.random.rand(3, 3) * 0.05 * i
+            atoms.cell = [10, 10, 10]
+            atoms.pbc = True
+            frames.append(atoms)
+        
+        write(traj_path, frames)
+        yield traj_path
+        shutil.rmtree(temp_dir)
+    
+    def test_msd_analysis_basic(self, test_trajectory_file):
+        """Test basic MSD analysis."""
+        result = msd_analysis(
+            test_trajectory_file,
+            timestep_fs=1.0,
+            indices_path=None  # Use correct parameter name
+        )
+        
+        # Should complete without error
+        assert result is not None
+    
+    def test_msd_analysis_with_indices_path(self, test_trajectory_file):
+        """Test analysis with indices_path parameter."""
+        # Create temp indices file
+        temp_dir = tempfile.mkdtemp()
+        indices_file = os.path.join(temp_dir, 'indices.npy')
+        np.save(indices_file, np.array([0, 1]))
+        
+        result = msd_analysis(
+            test_trajectory_file,
+            timestep_fs=1.0,
+            indices_path=indices_file
+        )
+        
+        assert result is not None
+        shutil.rmtree(temp_dir)
+
+
+class TestMSDParameterVariations:
+    """Test MSD calculations with various parameter combinations."""
+    
+    @pytest.fixture
+    def varied_trajectory(self):
+        """Create trajectory with varied properties."""
+        frames = []
+        for i in range(8):
+            atoms = Atoms('H4O2', positions=[
+                [0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0],
+                [0.5, 1, 0], [1.5, 1, 0]
+            ])
+            atoms.positions += np.random.rand(6, 3) * 0.02 * i
+            atoms.cell = [10, 10, 10]
+            frames.append(atoms)
+        return frames
+    
+    def test_msd_with_atom_subset_variations(self, varied_trajectory):
+        """Test MSD with different atom subsets."""
+        timestep = 1.0
+        
+        subsets = [
+            np.array([0]),
+            np.array([0, 1]),
+            np.array([0, 1, 2, 3]),
+            np.array([0, 1, 2, 3, 4, 5])
+        ]
+        
+        for atom_indices in subsets:
+            result = calculate_msd(varied_trajectory, timestep, atom_indices=atom_indices)
+            assert result is not None
+            msd_vals, msd_times = result
+            assert len(msd_vals) > 0
+    
+    def test_msd_with_different_timesteps(self, varied_trajectory):
+        """Test MSD with different timestep values."""
+        atom_indices = np.array([0, 1])
+        
+        for timestep in [0.5, 1.0, 2.0, 5.0]:
+            result = calculate_msd(varied_trajectory, timestep, atom_indices=atom_indices)
+            msd_vals, msd_times = result
+            
+            # Larger timestep should result in larger time values
+            assert msd_times[-1] > 0
+
+
+class TestMSDIntegration:
+    """Integration tests combining multiple MSD functions."""
+    
+    def test_full_msd_workflow(self):
+        """Test complete workflow from trajectory to diffusion coefficient."""
+        # Create trajectory
+        frames = []
+        for i in range(10):
+            atoms = Atoms('H2O', positions=[[0, 0, 0], [1, 0, 0], [0, 1, 0]])
+            atoms.positions += np.array([0.02*i, 0.02*i, 0])
+            frames.append(atoms)
+        
+        # Calculate MSD
+        timestep = 1.0
+        atom_indices = np.array([0, 1, 2])
+        msd_values, msd_times = calculate_msd(frames, timestep, atom_indices=atom_indices)
+        
+        # Just verify the MSD calculation worked
+        assert len(msd_values) > 0
+        assert len(msd_times) > 0
+        assert msd_values[0] >= 0
+    
+    def test_msd_with_different_atom_types(self):
+        """Test MSD calculation distinguishing atom types."""
+        frames = []
+        for i in range(8):
+            atoms = Atoms('H2O2', positions=[
+                [0, 0, 0], [1, 0, 0],  # H atoms
+                [0.5, 0.5, 0], [1.5, 0.5, 0]  # O atoms
+            ])
+            atoms.positions += np.random.rand(4, 3) * 0.01 * i
+            frames.append(atoms)
+        
+        timestep = 1.0
+        
+        # MSD for H atoms only
+        h_indices = np.array([0, 1])
+        msd_h, time_h = calculate_msd(frames, timestep, atom_indices=h_indices)
+        
+        # MSD for O atoms only
+        o_indices = np.array([2, 3])
+        msd_o, time_o = calculate_msd(frames, timestep, atom_indices=o_indices)
+        
+        assert len(msd_h) > 0
+        assert len(msd_o) > 0
+        # Both should have same time array
+        assert len(time_h) == len(time_o)
+
+
+class TestMSDEdgeCases:
+    """Test MSD with edge cases."""
+    
+    def test_msd_two_frame_trajectory(self):
+        """Test MSD with minimal 2-frame trajectory."""
+        frames = [
+            Atoms('H2', positions=[[0, 0, 0], [1, 0, 0]]),
+            Atoms('H2', positions=[[0.1, 0, 0], [1.1, 0, 0]])
+        ]
+        
+        result = calculate_msd(frames, 1.0, atom_indices=np.array([0, 1]))
+        assert result is not None
+    
+    def test_msd_large_trajectory(self):
+        """Test MSD with larger trajectory."""
+        frames = []
+        for i in range(50):
+            atoms = Atoms('H2O', positions=[[0, 0, 0], [1, 0, 0], [0, 1, 0]])
+            atoms.positions += np.random.rand(3, 3) * 0.01 * i
+            frames.append(atoms)
+        
+        result = calculate_msd(frames, 1.0, atom_indices=np.array([0, 1, 2]))
+        msd_vals, msd_times = result
+        
+        # MSD calculation may return N or N+1 points
+        assert len(msd_times) >= 50
+        assert len(msd_times) <= 51
+        # MSD values should increase over time (random walk)
+        assert msd_vals[0] == 0  # Initial MSD is zero