Coordinate Transformations

def translate(atomgroup, t):
    """
    Translate atomic coordinates by a vector.
    
    Parameters
    ----------
    atomgroup : AtomGroup
        Atoms to translate.
    t : array-like
        Translation vector as 3-element array [dx, dy, dz] in Angstrom.
        
    Examples
    --------
    >>> # Move protein 10 Å along x-axis
    >>> protein.translate([10.0, 0.0, 0.0])
    
    >>> # Center protein at origin
    >>> com = protein.center_of_mass()
    >>> protein.translate(-com)
    
    >>> # Move to specific position
    >>> target_position = [50.0, 25.0, 75.0]
    >>> current_com = protein.center_of_mass()
    >>> translation = target_position - current_com
    >>> protein.translate(translation)
    """

def rotate(atomgroup, R, point=(0, 0, 0)):
    """
    Rotate coordinates using a rotation matrix.
    
    Parameters
    ----------
    atomgroup : AtomGroup
        Atoms to rotate.
    R : array-like
        3x3 rotation matrix. Must be orthogonal with determinant +1.
    point : array-like, optional
        Point about which to rotate (default origin).
        
    Examples
    --------
    >>> import numpy as np
    >>> # 90-degree rotation around z-axis
    >>> R_z = np.array([[0, -1, 0],
    ...                 [1,  0, 0], 
    ...                 [0,  0, 1]])
    >>> protein.rotate(R_z)
    
    >>> # Rotate around protein center
    >>> center = protein.center_of_geometry()
    >>> protein.rotate(R_z, point=center)
    """

def rotateby(atomgroup, angle, axis, point=None):
    """
    Rotate by angle around axis.
    
    Parameters
    ----------
    atomgroup : AtomGroup
        Atoms to rotate.
    angle : float
        Rotation angle in degrees.
    axis : array-like
        Rotation axis as 3-element unit vector.
    point : array-like, optional
        Point about which to rotate (default group center).
        
    Examples
    --------
    >>> # Rotate 45 degrees around z-axis
    >>> protein.rotateby(45.0, [0, 0, 1])
    
    >>> # Rotate around custom axis
    >>> axis = [1, 1, 0] / np.sqrt(2)  # Normalize axis
    >>> protein.rotateby(90.0, axis)
    
    >>> # Rotate around specific point
    >>> ca_atom = protein.select_atoms("name CA")[0]
    >>> protein.rotateby(30.0, [0, 0, 1], point=ca_atom.position)
    """

def transform(atomgroup, M):
    """
    Apply transformation matrix to coordinates.
    
    Parameters
    ----------
    atomgroup : AtomGroup
        Atoms to transform.
    M : array-like
        Transformation matrix. Can be:
        - 3x3 rotation/scaling matrix
        - 4x4 homogeneous transformation matrix
        - 3x4 [rotation|translation] matrix
        
    Examples
    --------
    >>> # Combined rotation and translation using 4x4 matrix
    >>> M = np.eye(4)
    >>> M[:3, :3] = rotation_matrix  # 3x3 rotation
    >>> M[:3, 3] = translation_vector  # Translation
    >>> protein.transform(M)
    
    >>> # Scaling transformation
    >>> scale_matrix = np.diag([2.0, 1.0, 0.5])  # Scale x by 2, z by 0.5
    >>> protein.transform(scale_matrix)
    """

from MDAnalysis.analysis import align

def alignto(mobile, reference, select="all", weights=None, subsetA=None, 
            subsetB=None, **kwargs):
    """
    Align mobile structure to reference structure.
    
    Parameters
    ----------
    mobile : Universe or AtomGroup
        Structure to be aligned (modified in place).
    reference : Universe or AtomGroup  
        Reference structure for alignment.
    select : str, optional
        Selection string for atoms used in alignment (default "all").
    weights : str or array-like, optional
        Weights for alignment. Can be "mass" for mass weighting or
        array of weights matching selected atoms.
    subsetA : AtomGroup, optional
        Specific subset of mobile atoms for fitting.
    subsetB : AtomGroup, optional
        Specific subset of reference atoms for fitting.
    **kwargs
        Additional arguments for alignment algorithm.
        
    Returns
    -------
    dict
        Dictionary containing:
        - 'rmsd': RMSD after alignment
        - 'rotmat': 3x3 rotation matrix applied
        - 'translation': translation vector applied
        
    Examples
    --------
    >>> # Align protein backbones
    >>> result = align.alignto(mobile_u, reference_u, select="backbone")
    >>> print(f"Alignment RMSD: {result['rmsd']:.2f} Å")
    
    >>> # Mass-weighted alignment of CA atoms
    >>> result = align.alignto(mobile_u, reference_u, 
    ...                       select="name CA", weights="mass")
    
    >>> # Align using specific atom selections
    >>> mobile_ca = mobile_u.select_atoms("name CA")
    >>> ref_ca = reference_u.select_atoms("name CA") 
    >>> result = align.alignto(mobile_u, reference_u,
    ...                       subsetA=mobile_ca, subsetB=ref_ca)
    """

class AlignTraj:
    """
    Align trajectory frames to reference structure.
    
    Performs structural alignment of entire trajectory, optionally
    writing aligned frames to output file.
    """
    
    def __init__(self, mobile, reference, select="all", filename=None,
                 weights=None, in_memory=False, **kwargs):
        """
        Initialize trajectory alignment.
        
        Parameters
        ----------
        mobile : Universe
            Universe with trajectory to align.
        reference : Universe or AtomGroup
            Reference structure for alignment.
        select : str, optional
            Selection for alignment atoms (default "all").
        filename : str, optional
            Output filename for aligned trajectory.
        weights : str or array-like, optional
            Alignment weights ("mass" or custom array).
        in_memory : bool, optional
            Load trajectory in memory for speed (default False).
        **kwargs
            Additional alignment parameters.
            
        Examples
        --------
        >>> # Align entire trajectory and save
        >>> aligner = align.AlignTraj(u, reference, 
        ...                          select="protein and name CA",
        ...                          filename="aligned.xtc")
        >>> aligner.run()
        
        >>> # In-memory alignment for analysis
        >>> aligner = align.AlignTraj(u, reference, 
        ...                          select="backbone",
        ...                          in_memory=True)
        >>> aligner.run()
        >>> rmsd_data = aligner.rmsd  # RMSD per frame
        """
    
    def run(self, start=None, stop=None, step=None, **kwargs):
        """
        Execute trajectory alignment.
        
        Parameters
        ----------
        start : int, optional
            First frame to align (default None for beginning).
        stop : int, optional
            Last frame to align (default None for end).
        step : int, optional
            Step size for alignment (default None for 1).
        **kwargs
            Additional run parameters.
            
        Returns
        -------
        AlignTraj
            Self for method chaining.
        """
    
    @property
    def rmsd(self):
        """
        RMSD values for each aligned frame.
        
        Returns
        -------
        numpy.ndarray
            Array of RMSD values in Angstrom for each frame.
        """

from MDAnalysis.analysis.align import rotation_matrix

def rotation_matrix(a, b, weights=None):
    """
    Calculate optimal rotation matrix between two coordinate sets.
    
    Uses Kabsch algorithm to find rotation matrix that minimizes
    RMSD between coordinate sets.
    
    Parameters
    ----------
    a : array-like
        Coordinate set A of shape (n, 3).
    b : array-like  
        Coordinate set B of shape (n, 3).
    weights : array-like, optional
        Weights for each coordinate pair.
        
    Returns
    -------
    tuple
        (rotation_matrix, rmsd) where rotation_matrix is 3x3 array
        and rmsd is root mean square deviation after alignment.
        
    Examples
    --------
    >>> # Calculate rotation between CA atoms
    >>> mobile_ca = mobile.select_atoms("name CA").positions
    >>> ref_ca = reference.select_atoms("name CA").positions
    >>> R, rmsd = rotation_matrix(mobile_ca, ref_ca)
    >>> print(f"Optimal RMSD: {rmsd:.2f} Å")
    
    >>> # Apply rotation to mobile structure  
    >>> mobile.atoms.rotate(R)
    """

def fit_rot_trans(mobile, reference, weights=None):
    """
    Calculate optimal rotation and translation between structures.
    
    Parameters
    ----------
    mobile : array-like
        Mobile coordinates of shape (n, 3).
    reference : array-like
        Reference coordinates of shape (n, 3).
    weights : array-like, optional
        Weights for fitting.
        
    Returns
    -------
    tuple
        (rotation, translation, rmsd) where rotation is 3x3 matrix,
        translation is 3-element vector, and rmsd is final RMSD.
        
    Examples
    --------
    >>> R, t, rmsd = fit_rot_trans(mobile_coords, ref_coords)
    >>> # Apply transformation
    >>> mobile_coords_aligned = mobile_coords @ R.T + t
    """

def wrap(atomgroup, compound="atoms", center="com", box=None, inplace=True):
    """
    Wrap coordinates into primary unit cell.
    
    Parameters
    ----------
    atomgroup : AtomGroup
        Atoms to wrap.
    compound : str, optional
        Level for wrapping operation:
        - "atoms": wrap individual atoms
        - "group": wrap entire group as unit
        - "residues": wrap by residue  
        - "segments": wrap by segment
        - "molecules": wrap by molecule
        - "fragments": wrap by fragment
    center : str or array-like, optional
        Center for wrapping:
        - "com": center of mass
        - "cog": center of geometry  
        - array: specific center point
    box : array-like, optional
        Unit cell dimensions. If None, uses current trajectory box.
    inplace : bool, optional
        Whether to modify coordinates in place (default True).
        
    Returns
    -------
    numpy.ndarray or None
        Wrapped coordinates if inplace=False, otherwise None.
        
    Examples
    --------
    >>> # Wrap all atoms individually
    >>> u.atoms.wrap(compound="atoms")
    
    >>> # Wrap by residues to keep residues intact
    >>> u.atoms.wrap(compound="residues", center="com")
    
    >>> # Wrap around specific center
    >>> center_point = [25.0, 25.0, 25.0]
    >>> u.atoms.wrap(center=center_point)
    
    >>> # Get wrapped coordinates without modifying original
    >>> wrapped_coords = u.atoms.wrap(inplace=False)
    """

def unwrap(atomgroup, compound="fragments", reference="com", inplace=True):
    """
    Unwrap coordinates across periodic boundaries.
    
    Removes effects of periodic wrapping to make molecules whole
    and trajectories continuous.
    
    Parameters
    ----------
    atomgroup : AtomGroup
        Atoms to unwrap.
    compound : str, optional
        Level for unwrapping:
        - "atoms": unwrap individual atoms
        - "group": unwrap entire group
        - "residues": unwrap by residue
        - "segments": unwrap by segment  
        - "molecules": unwrap by molecule
        - "fragments": unwrap by fragment (default)
    reference : str or array-like, optional
        Reference for unwrapping:
        - "com": center of mass
        - "cog": center of geometry
        - array: specific reference point
    inplace : bool, optional
        Whether to modify coordinates in place (default True).
        
    Returns
    -------
    numpy.ndarray or None
        Unwrapped coordinates if inplace=False, otherwise None.
        
    Examples
    --------
    >>> # Unwrap molecular fragments  
    >>> u.atoms.unwrap(compound="fragments")
    
    >>> # Unwrap by residues
    >>> protein = u.select_atoms("protein")
    >>> protein.unwrap(compound="residues", reference="com")
    
    >>> # Unwrap entire trajectory
    >>> for ts in u.trajectory:
    ...     u.atoms.unwrap(compound="molecules")
    """

def pack_into_box(atomgroup, box=None, inplace=True):
    """
    Pack coordinates into simulation box.
    
    Ensures all coordinates are within box boundaries by
    applying appropriate translations.
    
    Parameters
    ----------
    atomgroup : AtomGroup
        Atoms to pack into box.
    box : array-like, optional
        Box dimensions as [a, b, c, alpha, beta, gamma].
        If None, uses current trajectory box.
    inplace : bool, optional
        Whether to modify coordinates in place (default True).
        
    Returns
    -------
    numpy.ndarray or None
        Packed coordinates if inplace=False, otherwise None.
        
    Examples
    --------
    >>> # Pack all atoms into current box
    >>> u.atoms.pack_into_box()
    
    >>> # Pack into custom box  
    >>> box_dims = [50.0, 50.0, 50.0, 90.0, 90.0, 90.0]
    >>> u.atoms.pack_into_box(box=box_dims)
    """

from MDAnalysis import transformations

# Available transformation classes
class TransformationBase:
    """
    Base class for trajectory transformations.
    
    Transformations can be applied during trajectory reading
    to modify coordinates on-the-fly.
    """
    
    def __call__(self, ts):
        """
        Apply transformation to timestep.
        
        Parameters
        ----------
        ts : Timestep
            Timestep object to transform.
            
        Returns
        -------
        Timestep
            Modified timestep.
        """

class unwrap(TransformationBase):
    """Unwrap coordinates transformation."""
    
    def __init__(self, atomgroup, **kwargs):
        """
        Initialize unwrapping transformation.
        
        Parameters
        ----------
        atomgroup : AtomGroup
            Atoms to unwrap in each frame.
        **kwargs
            Additional unwrapping parameters.
        """

class wrap(TransformationBase):
    """Wrap coordinates transformation."""
    
    def __init__(self, atomgroup, **kwargs):
        """
        Initialize wrapping transformation.
        
        Parameters
        ----------
        atomgroup : AtomGroup
            Atoms to wrap in each frame.
        **kwargs
            Additional wrapping parameters.
        """

class center_in_box(TransformationBase):
    """Center system in simulation box."""
    
    def __init__(self, atomgroup, center='mass', **kwargs):
        """
        Initialize centering transformation.
        
        Parameters
        ----------
        atomgroup : AtomGroup
            Atoms to center.
        center : str, optional
            Centering method ('mass' or 'geometry').
        """

class fit_rot_trans(TransformationBase):
    """Align trajectory to reference structure."""
    
    def __init__(self, atomgroup, reference, **kwargs):
        """
        Initialize alignment transformation.
        
        Parameters
        ----------
        atomgroup : AtomGroup
            Atoms to align in each frame.
        reference : AtomGroup or array-like
            Reference coordinates for alignment.
        """

# Usage with Universe
transformations_list = [
    transformations.unwrap(u.atoms),
    transformations.center_in_box(u.select_atoms("protein")),
    transformations.wrap(u.atoms, compound="residues")
]

u = mda.Universe("topology.psf", "trajectory.dcd", 
                transformations=transformations_list)

# Now trajectory is automatically transformed during reading
for ts in u.trajectory:
    # Coordinates are already unwrapped, centered, and wrapped
    pass

def align_principal_axis(atomgroup, axis='z', vector=None):
    """
    Align principal axis of molecule with coordinate axis.
    
    Parameters
    ----------
    atomgroup : AtomGroup
        Atoms to align.
    axis : str or int, optional
        Target coordinate axis ('x', 'y', 'z' or 0, 1, 2).
    vector : array-like, optional
        Specific vector to align with (overrides axis).
        
    Examples
    --------
    >>> # Align protein's principal axis with z-axis
    >>> protein = u.select_atoms("protein")
    >>> align_principal_axis(protein, axis='z')
    
    >>> # Align with custom vector
    >>> target_vector = [1, 1, 0] / np.sqrt(2)
    >>> align_principal_axis(protein, vector=target_vector)
    """
    # Calculate principal axes
    principal_axes = atomgroup.principal_axes()
    longest_axis = principal_axes[0]  # Largest eigenvalue
    
    if vector is None:
        # Convert axis specification to vector
        if axis == 'x' or axis == 0:
            target = np.array([1, 0, 0])
        elif axis == 'y' or axis == 1: 
            target = np.array([0, 1, 0])
        elif axis == 'z' or axis == 2:
            target = np.array([0, 0, 1])
        else:
            raise ValueError("axis must be 'x', 'y', 'z' or 0, 1, 2")
    else:
        target = np.array(vector)
        target = target / np.linalg.norm(target)
    
    # Calculate rotation matrix
    from MDAnalysis.lib.mdamath import normal
    rotation_axis = np.cross(longest_axis, target)
    
    if np.allclose(rotation_axis, 0):
        # Axes are already aligned or antiparallel
        if np.dot(longest_axis, target) < 0:
            # Antiparallel - rotate 180 degrees
            perpendicular = normal(longest_axis)
            angle = np.pi
        else:
            # Already aligned
            return
    else:
        rotation_axis = rotation_axis / np.linalg.norm(rotation_axis)
        angle = np.arccos(np.clip(np.dot(longest_axis, target), -1, 1))
    
    # Create rotation matrix using Rodrigues' formula
    cos_angle = np.cos(angle)
    sin_angle = np.sin(angle)
    
    # Cross product matrix
    K = np.array([[0, -rotation_axis[2], rotation_axis[1]],
                  [rotation_axis[2], 0, -rotation_axis[0]], 
                  [-rotation_axis[1], rotation_axis[0], 0]])
    
    R = np.eye(3) + sin_angle * K + (1 - cos_angle) * np.dot(K, K)
    
    # Apply rotation
    center = atomgroup.center_of_geometry()
    atomgroup.translate(-center)
    atomgroup.rotate(R)
    atomgroup.translate(center)

def cartesian_to_spherical(coordinates, origin=None):
    """
    Convert Cartesian to spherical coordinates.
    
    Parameters
    ----------
    coordinates : array-like
        Cartesian coordinates of shape (..., 3).
    origin : array-like, optional
        Origin for coordinate system (default [0, 0, 0]).
        
    Returns
    -------
    numpy.ndarray
        Spherical coordinates [r, theta, phi] where:
        - r: radial distance
        - theta: polar angle (0 to π)  
        - phi: azimuthal angle (0 to 2π)
        
    Examples
    --------
    >>> coords = protein.positions
    >>> com = protein.center_of_mass()
    >>> spherical = cartesian_to_spherical(coords, origin=com)
    >>> r = spherical[..., 0]     # Distances from center
    >>> theta = spherical[..., 1]  # Polar angles
    >>> phi = spherical[..., 2]    # Azimuthal angles
    """
    coords = np.asarray(coordinates)
    if origin is not None:
        coords = coords - np.asarray(origin)
    
    # Radial distance
    r = np.sqrt(np.sum(coords**2, axis=-1))
    
    # Polar angle (theta)
    theta = np.arccos(coords[..., 2] / r)
    
    # Azimuthal angle (phi)  
    phi = np.arctan2(coords[..., 1], coords[..., 0])
    phi = np.where(phi < 0, phi + 2*np.pi, phi)  # Ensure 0 to 2π
    
    return np.stack([r, theta, phi], axis=-1)

def cylindrical_coordinates(coordinates, axis='z', origin=None):
    """
    Convert to cylindrical coordinates.
    
    Parameters
    ----------
    coordinates : array-like
        Cartesian coordinates of shape (..., 3).
    axis : str or int, optional
        Cylinder axis ('x', 'y', 'z' or 0, 1, 2).
    origin : array-like, optional
        Origin for coordinate system.
        
    Returns
    -------
    numpy.ndarray
        Cylindrical coordinates [rho, phi, z] where:
        - rho: radial distance from axis
        - phi: azimuthal angle  
        - z: position along axis
        
    Examples
    --------
    >>> # Analyze protein in cylindrical coordinates around z-axis
    >>> coords = protein.positions
    >>> cylindrical = cylindrical_coordinates(coords, axis='z')
    >>> rho = cylindrical[..., 0]  # Distance from z-axis
    >>> phi = cylindrical[..., 1]  # Angle around z-axis
    >>> z = cylindrical[..., 2]    # Height along z-axis
    """
    coords = np.asarray(coordinates)
    if origin is not None:
        coords = coords - np.asarray(origin)
    
    # Determine axis index
    if axis == 'x' or axis == 0:
        ax_idx = 0
        perp_indices = [1, 2]
    elif axis == 'y' or axis == 1:
        ax_idx = 1  
        perp_indices = [0, 2]
    elif axis == 'z' or axis == 2:
        ax_idx = 2
        perp_indices = [0, 1]
    else:
        raise ValueError("axis must be 'x', 'y', 'z' or 0, 1, 2")
    
    # Axial coordinate
    z = coords[..., ax_idx]
    
    # Radial distance in perpendicular plane
    rho = np.sqrt(coords[..., perp_indices[0]]**2 + 
                  coords[..., perp_indices[1]]**2)
    
    # Azimuthal angle
    phi = np.arctan2(coords[..., perp_indices[1]], 
                     coords[..., perp_indices[0]])
    phi = np.where(phi < 0, phi + 2*np.pi, phi)
    
    return np.stack([rho, phi, z], axis=-1)

def alignment_workflow(mobile_universe, reference_universe):
    """
    Complete structural alignment workflow example.
    """
    # Step 1: Initial alignment using backbone atoms
    print("Performing initial backbone alignment...")
    result = align.alignto(mobile_universe, reference_universe, 
                          select="backbone")
    print(f"Initial RMSD: {result['rmsd']:.2f} Å")
    
    # Step 2: Refined alignment using CA atoms with mass weighting
    print("Refining alignment with CA atoms...")
    result = align.alignto(mobile_universe, reference_universe,
                          select="name CA", weights="mass")
    print(f"Refined RMSD: {result['rmsd']:.2f} Å")
    
    # Step 3: Calculate RMSD for different regions
    regions = {
        'backbone': 'backbone',
        'sidechains': 'protein and not backbone',
        'alpha_carbons': 'name CA',
        'binding_site': 'resid 23 45 67 89'
    }
    
    print("\nRegion-specific RMSDs after alignment:")
    for region_name, selection in regions.items():
        mobile_sel = mobile_universe.select_atoms(selection)
        ref_sel = reference_universe.select_atoms(selection)
        
        if len(mobile_sel) > 0 and len(ref_sel) > 0:
            # Calculate RMSD manually
            diff = mobile_sel.positions - ref_sel.positions
            rmsd = np.sqrt(np.mean(np.sum(diff**2, axis=1)))
            print(f"  {region_name}: {rmsd:.2f} Å")
    
    return result

# Run alignment workflow
result = alignment_workflow(mobile_u, reference_u)

def preprocess_trajectory(universe, output_file):
    """
    Comprehensive trajectory preprocessing pipeline.
    """
    protein = universe.select_atoms("protein")
    
    # Define transformation pipeline  
    transformations_list = [
        # 1. Unwrap molecules to make them whole
        transformations.unwrap(universe.atoms, compound="molecules"),
        
        # 2. Center protein in box
        transformations.center_in_box(protein, center='mass'),
        
        # 3. Align protein backbone to reference (first frame)
        transformations.fit_rot_trans(
            protein.select_atoms("backbone"),
            protein.select_atoms("backbone").positions  # First frame as reference
        ),
        
        # 4. Wrap solvent around protein
        transformations.wrap(universe.atoms, compound="residues")
    ]
    
    # Apply transformations and write processed trajectory
    with mda.Writer(output_file, n_atoms=universe.atoms.n_atoms) as W:
        for ts in universe.trajectory:
            # Apply transformations
            for transform in transformations_list:
                ts = transform(ts)
            
            # Write processed frame
            W.write(universe.atoms, ts=ts)
    
    print(f"Processed trajectory written to: {output_file}")

# Process trajectory
preprocess_trajectory(u, "processed_trajectory.xtc")

def analyze_pbc_effects(universe):
    """
    Analyze effects of periodic boundary conditions.
    """
    protein = universe.select_atoms("protein")
    
    results = {
        'frame': [],
        'original_rgyr': [],
        'wrapped_rgyr': [], 
        'unwrapped_rgyr': [],
        'box_violations': []
    }
    
    for ts in universe.trajectory:
        frame = ts.frame
        box_dims = ts.dimensions[:3]  # Box lengths
        
        # Original coordinates
        original_rgyr = protein.radius_of_gyration()
        
        # Wrapped coordinates
        protein_wrapped = protein.copy()
        protein_wrapped.wrap(compound="residues")
        wrapped_rgyr = protein_wrapped.radius_of_gyration()
        
        # Unwrapped coordinates
        protein_unwrapped = protein.copy()
        protein_unwrapped.unwrap(compound="residues")
        unwrapped_rgyr = protein_unwrapped.radius_of_gyration()
        
        # Check for atoms outside box
        coords = protein.positions
        violations = 0
        for dim in range(3):
            if box_dims[dim] > 0:  # Valid box dimension
                outside = np.sum((coords[:, dim] < 0) | 
                               (coords[:, dim] > box_dims[dim]))
                violations += outside
        
        # Store results
        results['frame'].append(frame)
        results['original_rgyr'].append(original_rgyr)
        results['wrapped_rgyr'].append(wrapped_rgyr)
        results['unwrapped_rgyr'].append(unwrapped_rgyr) 
        results['box_violations'].append(violations)
    
    # Analysis summary
    print("Periodic Boundary Condition Analysis:")
    print(f"Average Rgyr (original): {np.mean(results['original_rgyr']):.2f} Å")
    print(f"Average Rgyr (wrapped): {np.mean(results['wrapped_rgyr']):.2f} Å")
    print(f"Average Rgyr (unwrapped): {np.mean(results['unwrapped_rgyr']):.2f} Å")
    print(f"Frames with box violations: {np.sum(np.array(results['box_violations']) > 0)}")
    
    return results

# Analyze PBC effects
pbc_analysis = analyze_pbc_effects(u)

class RemoveCOMMotion(transformations.TransformationBase):
    """
    Custom transformation to remove center-of-mass motion.
    
    Removes translational motion by centering system at origin
    and removes rotational motion by aligning to reference.
    """
    
    def __init__(self, atomgroup, reference_frame=0):
        """
        Initialize COM motion removal.
        
        Parameters
        ----------
        atomgroup : AtomGroup
            Atoms for COM calculation and motion removal.
        reference_frame : int, optional
            Frame to use as reference for rotation removal.
        """
        self.atomgroup = atomgroup
        self.reference_frame = reference_frame
        self.reference_coords = None
        self._first_frame = True
    
    def __call__(self, ts):
        """Apply transformation to timestep."""
        # Get current frame coordinates
        current_coords = self.atomgroup.positions.copy()
        
        # Remove translational motion (center at origin)
        com = np.mean(current_coords, axis=0)
        current_coords -= com
        
        # Set up reference for rotational alignment
        if self._first_frame and ts.frame == self.reference_frame:
            self.reference_coords = current_coords.copy()
            self._first_frame = False
        
        # Remove rotational motion (align to reference)
        if self.reference_coords is not None:
            from MDAnalysis.analysis.align import rotation_matrix
            R, _ = rotation_matrix(current_coords, self.reference_coords)
            current_coords = current_coords @ R.T
        
        # Update positions in timestep
        self.atomgroup.positions = current_coords
        
        return ts

# Use custom transformation
protein = u.select_atoms("protein and name CA")
custom_transform = RemoveCOMMotion(protein, reference_frame=0)

# Apply during trajectory reading
u_processed = mda.Universe("topology.psf", "trajectory.dcd",
                          transformations=[custom_transform])

for ts in u_processed.trajectory:
    # Protein motion has been removed
    pass