Examples
========

This page provides examples of using TACO Format for various common tasks.

Basic Python Examples
---------------------

Here are some basic examples from the test suite showing how to use TACO Format with Python.

Writing and Reading a Single Structure
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

.. code-block:: python

    import taco_format
    import ase
    import numpy as np
    import tempfile
    import os
    
    # Create a temporary file
    with tempfile.NamedTemporaryFile(suffix='.taco', delete=False) as tmp:
        filepath = tmp.name
    
    try:
        # Create a simple ASE Atoms object
        atoms = ase.Atoms('H2O', 
                          positions=[[0.0, 0.0, 0.0], 
                                     [0.0, 0.0, 0.95], 
                                     [0.0, 0.9, 0.0]])
        
        # Save using taco_format
        taco_format.save_ase(filepath, atoms)
        
        # Load it back
        loaded_atoms = taco_format.read(filepath)[0]
        
        # Check that the data is preserved
        assert len(loaded_atoms) == len(atoms)
        np.testing.assert_allclose(loaded_atoms.positions, atoms.positions)
    finally:
        # Clean up
        os.unlink(filepath)

Working with Trajectories
~~~~~~~~~~~~~~~~~~~~~~~~~

.. code-block:: python

    import taco_format
    import ase
    import numpy as np
    import tempfile
    import os
    
    # Create a temporary file
    with tempfile.NamedTemporaryFile(suffix='.taco', delete=False) as tmp:
        filepath = tmp.name
    
    try:
        # Create a simple trajectory with 5 frames
        trajectory = []
        for i in range(5):
            # Create atoms with slightly different positions in each frame
            atoms = ase.Atoms('H2O',
                              positions=[[i*0.01, 0.0, 0.0],
                                         [i*0.01, 0.0, 0.95],
                                         [i*0.01, 0.9, 0.0]])
            trajectory.append(atoms)
            
        # Write trajectory to TACO file
        taco_format.write(filepath, trajectory)
        
        # Read the entire trajectory back
        loaded_trajectory = taco_format.read(filepath)
        
        # Check number of frames
        assert len(loaded_trajectory) == len(trajectory)
        
        # Read specific frame (frame 2)
        specific_frame = taco_format.read(filepath, frame_index=2)[0]
        np.testing.assert_allclose(specific_frame.positions, trajectory[2].positions)
        
        # Read a range of frames (frames 1-3)
        frame_range = taco_format.read_frame_range(filepath, 1, 4)
        assert len(frame_range) == 3  # Should return frames 1, 2, and 3
    finally:
        # Clean up
        os.unlink(filepath)

Setting Compression Options
~~~~~~~~~~~~~~~~~~~~~~~~~~~

.. code-block:: python

    import taco_format
    import ase
    import os
    
    # Create a water molecule
    atoms = ase.Atoms('H2O', 
                      positions=[[0.0, 0.0, 0.0], 
                                 [0.0, 0.0, 0.95], 
                                 [0.0, 0.9, 0.0]])
    
    # Save with default compression settings
    taco_format.save_ase('water_default.taco', atoms)
    
    # Save with lossless compression
    taco_format.save_ase('water_lossless.taco', atoms, lossless=True)
    
    # Save with high compression level
    taco_format.save_ase('water_high_compression.taco', atoms, compression_level=18)
    
    # Compare file sizes
    size_default = os.path.getsize('water_default.taco')
    size_lossless = os.path.getsize('water_lossless.taco')
    size_high_comp = os.path.getsize('water_high_compression.taco')
    
    print(f"Default compression: {size_default} bytes")
    print(f"Lossless compression: {size_lossless} bytes")
    print(f"High compression: {size_high_comp} bytes")
    
    # Clean up
    for filename in ['water_default.taco', 'water_lossless.taco', 'water_high_compression.taco']:
        if os.path.exists(filename):
            os.unlink(filename)

Converting Between Formats
--------------------------

Converting from XYZ to TACO
~~~~~~~~~~~~~~~~~~~~~~~~~~~

.. code-block:: python

    import ase.io
    import taco_format
    
    # Read XYZ file
    atoms_list = ase.io.read('trajectory.xyz', index=':')
    
    # Save as TACO
    taco_format.write('trajectory.taco', atoms_list, 
                      time_step=0.001,  # Time step in ps
                      full_frame_interval=10)  # Store full frame every 10 frames

Converting from TACO to XYZ
~~~~~~~~~~~~~~~~~~~~~~~~~~~

.. code-block:: python

    import ase.io
    import taco_format
    
    # Read TACO file
    atoms_list = taco_format.read('trajectory.taco')
    
    # Write as XYZ
    ase.io.write('trajectory.xyz', atoms_list)

Extracting Subsets of a Trajectory
----------------------------------

Extract Specific Frames
~~~~~~~~~~~~~~~~~~~~~~~

.. code-block:: python

    import taco_format
    
    # Extract frames 10, 20, 30, etc.
    selected_frames = taco_format.read_frames('trajectory.taco', 
                                           [10, 20, 30, 40, 50])
    
    # Save selected frames to a new file
    taco_format.write('selected_frames.taco', selected_frames)

Extract a Subset of Atoms
~~~~~~~~~~~~~~~~~~~~~~~~~

.. code-block:: python

    import taco_format
    
    # Extract atoms 0, 1, and 5 from all frames
    taco_format.extract_atoms(
        src_path='trajectory.taco',
        dst_path='subset.taco',
        atom_indices=[0, 1, 5]
    )

Analysis Examples
-----------------

Calculate RMSD Between Structures
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

.. code-block:: python

    import taco_format
    import numpy as np
    
    # Load two structures
    atoms1 = taco_format.read('structure1.taco')[0]
    atoms2 = taco_format.read('structure2.taco')[0]
    
    # Get positions
    pos1 = atoms1.get_positions()
    pos2 = atoms2.get_positions()
    
    # Calculate RMSD
    rmsd = taco_format.calc_rmsd(pos1, pos2)
    print(f"RMSD between structures: {rmsd:.6f} Å")

Calculate Center of Mass
~~~~~~~~~~~~~~~~~~~~~~~~

.. code-block:: python

    import taco_format
    import numpy as np
    
    # Load structure
    atoms = taco_format.read('structure.taco')[0]
    
    # Get positions
    positions = atoms.get_positions()
    
    # Get masses
    masses = atoms.get_masses()
    
    # Calculate center of mass
    com = taco_format.center_of_mass(positions, masses)
    print(f"Center of mass: {com}")

Working with Large Trajectories
-------------------------------

Memory-Efficient Reading
~~~~~~~~~~~~~~~~~~~~~~~~

For very large trajectories, you may want to process frames one by one:

.. code-block:: python

    import taco_format
    import os
    
    # Get file info
    file_info = taco_format.get_file_info('large_trajectory.taco')
    print(file_info)
    
    # Process frames one by one
    for i in range(10):  # Process first 10 frames
        frame = taco_format.read('large_trajectory.taco', frame_index=i)[0]
        
        # Do some analysis with each frame
        print(f"Frame {i}: potential energy = {frame.get_potential_energy()} eV")

Rust Examples
-------------

Here are examples of using TACO Format with Rust.

Writing and Reading a Simple Structure
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

.. code-block:: rust

    use taco_format::{Writer, Frame, FrameData};
    use ndarray::Array2;
    use tempfile::tempdir;
    use std::path::PathBuf;
    
    // Create a temporary directory for our test file
    let dir = tempdir().unwrap();
    let file_path = dir.path().join("test.taco");
    
    // Create test positions
    let positions = Array2::from_shape_fn((10, 3), |(i, j)| i as f32 + j as f32);
    
    // Create a frame with positions
    let frame_data = FrameData::new(positions);
    let frame = Frame::new(0, 0.0, frame_data);
    
    // Write the frame to a file
    let mut writer = Writer::create(&file_path).unwrap();
    writer.write_frame(frame).unwrap();
    writer.finish().unwrap();
    
    // Read the frame back
    let mut reader = taco_format::Reader::open(&file_path).unwrap();
    let read_frame = reader.read_frame(0).unwrap();
    
    // Access the positions
    let read_positions = read_frame.data.positions.unwrap();
    assert_eq!(read_positions.shape(), &[10, 3]);

Working with Trajectories in Rust
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

.. code-block:: rust

    use taco_format::{Writer, Frame, FrameData};
    use ndarray::Array2;
    use tempfile::tempdir;
    
    // Create a temporary directory for our test file
    let dir = tempdir().unwrap();
    let file_path = dir.path().join("trajectory.taco");
    
    // Create a writer
    let mut writer = Writer::create(&file_path).unwrap();
    
    // Write multiple frames with slightly changing positions
    let num_frames = 5;
    for i in 0..num_frames {
        let positions = Array2::from_shape_fn((10, 3), |(a, b)| {
            (a as f32 + b as f32) + (i as f32 * 0.1)
        });
        
        let frame_data = FrameData::new(positions);
        let frame = Frame::new(i as u64, i as f64 * 0.001, frame_data);
        
        writer.write_frame(frame).unwrap();
    }
    
    writer.finish().unwrap();
    
    // Read frames back
    let mut reader = taco_format::Reader::open(&file_path).unwrap();
    
    // Check number of frames
    assert_eq!(reader.num_frames(), num_frames);
    
    // Read specific frame
    let frame_idx = 2;
    let frame = reader.read_frame(frame_idx).unwrap();
    
    // Check frame metadata
    assert_eq!(frame.frame_number, frame_idx as u64);
    assert_eq!(frame.time, (frame_idx as f64) * 0.001);

Creating a TACO File From Scratch in Rust
-----------------------------------------

.. code-block:: rust

    use taco_format::{Writer, Header, Frame, FrameData, SimulationMetadata, AtomMetadata};
    use ndarray::Array2;
    
    // Create positions for 100 atoms
    let mut positions = Array2::<f32>::zeros((100, 3));
    
    // Fill with some data
    for i in 0..100 {
        positions[[i, 0]] = i as f32 * 0.1;
        positions[[i, 1]] = i as f32 * 0.2;
        positions[[i, 2]] = i as f32 * 0.3;
    }
    
    // Create a default header
    let header = Header::new(
        100, // Number of atoms
        0.001, // Time step in ps
        SimulationMetadata::default(),
        AtomMetadata::default(),
        Default::default(), // Default compression settings
    );
    
    // Create a writer with our header
    let mut writer = Writer::create_with_header("trajectory.taco", header).unwrap();
    
    // Create and write frames
    for i in 0..10 {
        // Modify positions slightly for each frame
        for j in 0..100 {
            positions[[j, 0]] += 0.01;
        }
        
        // Create frame data with positions
        let frame_data = FrameData::new(positions.clone());
        
        // Create a frame
        let frame = Frame::new(
            i as u64, // Frame number
            i as f64 * 0.001, // Time
            frame_data,
        );
        
        // Write the frame
        writer.write_frame(frame).unwrap();
    }
    
    // Finish writing
    writer.finish().unwrap();