Preserving Clusters in Error-Bounded Lossy Compression of Particle Data

1. Problem Statement

Large-scale scientific simulations (e.g., cosmology, molecular dynamics, fluid dynamics) generate massive particle datasets that require error-bounded lossy compression to manage storage and I/O costs. While existing compressors (like SZ3, ZFP, Draco) guarantee that the difference between original and decompressed coordinates remains within a user-defined error bound ($\xi$), they fail to preserve topological structures, specifically single-linkage clusters (Friends-of-Friends or FoF).

• The Core Issue: Single-linkage clustering relies on a strict distance threshold ($b$). Even small coordinate perturbations caused by compression can:
  • Break valid links: Particles within distance $b$ are pushed apart, fragmenting a single cluster (halo) into multiple pieces.
  • Create false links: Particles originally separated by $>b$ are pushed together, artificially merging distinct clusters.
• Current Limitations:
  • Strict Error Bounds: To preserve clusters, users must set extremely tight error bounds (e.g., $10^{-8}$), which drastically reduces compression ratios (often to near-lossless levels, e.g., 1.4x–1.7x).
  • Auxiliary Storage: Storing cluster assignments separately incurs significant overhead (up to 10% of data size) and requires re-computation for every new parameter set.
  • Lack of Guarantees: Standard compressors offer no guarantee that the scientific validity of downstream analyses (like Halo Mass Functions) is maintained.

2. Methodology

The authors propose a correction-based technique applied in situ after decompression. The method treats cluster preservation as a constrained optimization problem.

A. Formulation

• Vulnerable Pairs: The algorithm identifies "vulnerable" particle pairs where the original distance $d_{i,j}$ is within a critical range of the linking length $b$ (specifically $(b - 2\sqrt{3}\xi, b + 2\sqrt{3}\xi]$). Only these pairs are susceptible to connectivity changes.
• Editable Particles: Particles involved in vulnerable pairs are designated as "editable."
• Optimization Objective: Minimize a loss function $L$ that penalizes violations of the linking constraint:
  • Broken Links: If $d_{i,j} \le b$ but decompressed $\hat{d}_{i,j} > b$, penalize the overflow.
  • False Links: If $d_{i,j} > b$ but decompressed $\hat{d}_{i,j} \le b$, penalize the intrusion.
• Constraints: The corrected coordinates must strictly satisfy the original error bound $|\hat{x} - x| \le \xi$.

B. Algorithm: Projected Gradient Descent (PGD)

1. Detection: Uses spatial partitioning (cell-linking) to efficiently find vulnerable pairs in $O(N \cdot \bar{\rho})$ time.
2. Optimization: Iteratively updates the positions of editable particles using gradient descent to minimize $L$.
3. Projection: After each gradient step, coordinates are projected back onto the "box" defined by the error bound $\xi$ to ensure constraints are met.
4. Safety Margin: To account for quantization errors during the final encoding of edits, the algorithm optimizes a "tightened" loss function ($L_{tight}$) with a slightly stricter threshold, ensuring robustness after bit-depth reduction.
5. Encoding: The coordinate shifts (edits) are quantized, packed into a bitmask, and compressed losslessly (Huffman + ZSTD) to create a compact correction stream.

C. Implementation & Scalability

• GPU Acceleration: The entire pipeline (spatial partitioning, pair detection, PGD updates) is parallelized on GPUs.
  • Uses atomic operations and shared-memory reductions for efficient gradient accumulation.
  • Achieves up to 62× speedup over CPU baselines.
• Distributed Memory (MPI): Scales to multi-node clusters by exchanging "ghost zones" (boundary particles) between ranks. The ghost zone width is $\delta = b + 2\sqrt{3}\xi$ to ensure all cross-boundary links are detected.

3. Key Contributions

1. Novel Correction Algorithm: A method to preserve single-linkage clusters under lossy compression by formulating it as a constrained optimization problem solved via PGD, operating on decompressed data from off-the-shelf compressors.
2. Scalable Parallel Implementation: A distributed, data-parallel solver supporting both CPU and GPU backends, capable of handling exascale datasets (tens of trillions of particles) via MPI.
3. Comprehensive Evaluation: Rigorous validation across cosmology (HACC), molecular dynamics (EXAALT), and fluid dynamics (FPM) datasets, demonstrating superior cluster preservation compared to state-of-the-art compressors.

4. Experimental Results

The method was evaluated on datasets ranging from 200 MB to 192 GB using compressors like SZ3, ZFP, Draco, and LCP.

• Cluster Preservation (Accuracy):
  • Matthews Correlation Coefficient (MCC): The proposed method achieves near-perfect MCC (close to 1.0) for vulnerable pairs, whereas base compressors often drop significantly (e.g., <0.5 for SZ3 on HACC).
  • Halo Mass Function (HMF): The statistical distribution of halo masses is preserved with significantly lower relative error compared to base compressors, ensuring scientific validity.
• Compression Efficiency:
  • The method introduces negligible storage overhead for datasets with low violation density (e.g., FPM fluid dynamics).
  • For high-density datasets, the overhead is manageable, often resulting in a net gain in effective compression ratio compared to using ultra-strict error bounds to force preservation.
• Performance:
  • Throughput: The GPU implementation achieves 62× speedup over a 64-core CPU baseline.
  • Scalability: Strong scaling tests on 256 MPI ranks show the GPU version maintains high efficiency, though communication overhead becomes a bottleneck at extreme scales. Weak scaling tests demonstrate near-ideal throughput scaling up to 1,024 ranks.
• Convergence: The algorithm typically converges to zero loss (perfect cluster restoration) within a few iterations for moderate error bounds. Even with a truncated iteration limit (e.g., 100 iterations), it significantly outperforms base compressors.

5. Significance

This work bridges a critical gap in scientific data management. It allows researchers to utilize high-ratio lossy compression without sacrificing the integrity of topological analyses (clustering), which are fundamental to understanding the universe's structure, material properties, and fluid behaviors.

• Scientific Impact: Enables the storage and analysis of exascale particle datasets that would otherwise be impossible to store due to I/O and capacity limits, while guaranteeing that derived scientific metrics (like halo mass functions) remain accurate.
• Technical Innovation: Demonstrates that feature-preserving compression can be achieved as a post-processing optimization step, decoupling the need for complex, feature-aware base compressors.
• Practicality: The method is compatible with existing workflows (works with SZ3, Draco, etc.) and integrates seamlessly into high-throughput simulation pipelines via GPU acceleration.

Limitations & Future Work:

• The correction is specific to the linking length $b$ used at compression time; changing $b$ requires re-correction.
• Currently assumes non-periodic domains (though the paper notes it can be extended).
• Future work aims to integrate correction directly into the compressor's transform stage and extend to phase-space halo finders (e.g., Rockstar).