diff --git a/src/main/java/org/apache/sysds/api/DMLScript.java b/src/main/java/org/apache/sysds/api/DMLScript.java index c286b8d3b52..8391634fdb1 100644 --- a/src/main/java/org/apache/sysds/api/DMLScript.java +++ b/src/main/java/org/apache/sysds/api/DMLScript.java @@ -131,6 +131,8 @@ public class DMLScript public static double GPU_MEMORY_UTILIZATION_FACTOR = 0.9; // Set GPU memory allocator to use public static String GPU_MEMORY_ALLOCATOR = "cuda"; + // Enable/disable federated compression + public static boolean FEDERATED_COMPRESSION = false; // Enable/disable lineage tracing public static boolean LINEAGE = DMLOptions.defaultOptions.lineage; // Enable/disable deduplicate lineage items diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederationMap.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederationMap.java index 68dff4785b3..4d7a06e5614 100644 --- a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederationMap.java +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederationMap.java @@ -42,12 +42,19 @@ import org.apache.sysds.hops.fedplanner.FTypes.AlignType; import org.apache.sysds.hops.fedplanner.FTypes.FType; import org.apache.sysds.lops.RightIndex; +import org.apache.sysds.api.DMLScript; import org.apache.sysds.runtime.DMLRuntimeException; import org.apache.sysds.runtime.controlprogram.caching.CacheBlock; import org.apache.sysds.runtime.controlprogram.caching.CacheableData; import org.apache.sysds.runtime.controlprogram.federated.FederatedRequest.RequestType; import org.apache.sysds.runtime.instructions.cp.CPOperand; import org.apache.sysds.runtime.instructions.cp.ScalarObject; +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressionConfig; +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressionFactory; +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressionType; +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressedMatrix; +import org.apache.sysds.runtime.controlprogram.federated.compression.MatrixCompressor; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; import org.apache.sysds.runtime.instructions.cp.VariableCPInstruction; import org.apache.sysds.runtime.lineage.LineageItem; import org.apache.sysds.runtime.util.IndexRange; @@ -138,6 +145,27 @@ private FederatedRequest broadcast(CacheableData data, LineageItem lineageIte // is fine, because with broadcast all data on all workers) data.setFedMapping(copyWithNewIDAndRange( cb.getNumRows(), cb.getNumColumns(), id, FType.BROADCAST)); + + // === COMPRESSION INTEGRATION === + // Attempt TopK compression if the block is a MatrixBlock + if(DMLScript.FEDERATED_COMPRESSION && cb instanceof MatrixBlock) { + try { + CompressionConfig config = CompressionConfig.builder() + .enable(true) + .withType(CompressionType.TOPK) + .withSparsity(0.01) + .build(); + MatrixCompressor compressor = CompressionFactory.create(config); + CompressedMatrix compressed = compressor.compress((MatrixBlock) cb); + MatrixBlock decompressed = compressor.decompress(compressed); + return new FederatedRequest(RequestType.PUT_VAR, lineageItem, id, decompressed); + } + catch(Exception ex) { + // Fall back to uncompressed on any error + } + } + // === END COMPRESSION INTEGRATION === + return new FederatedRequest(RequestType.PUT_VAR, lineageItem, id, cb); } diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressedMatrix.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressedMatrix.java new file mode 100644 index 00000000000..4ad1660bfa2 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressedMatrix.java @@ -0,0 +1,81 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.controlprogram.federated.compression; + +import java.io.Serializable; + +/** + * Generic container for compressed matrix data. + * Stores the compressed representation along with metadata + * needed for decompression and size estimation. + */ +public class CompressedMatrix implements Serializable { + + private static final long serialVersionUID = 1L; + + private final CompressionType type; + private final int numRows; + private final int numCols; + private final Object compressedData; // Technique-specific data + private final double compressionRatio; + private final byte[] metadata; // Optional: scaling factors, etc. + + public CompressedMatrix(CompressionType type, int numRows, int numCols, + Object compressedData, double compressionRatio) { + this(type, numRows, numCols, compressedData, compressionRatio, null); + } + + public CompressedMatrix(CompressionType type, int numRows, int numCols, + Object compressedData, double compressionRatio, + byte[] metadata) { + this.type = type; + this.numRows = numRows; + this.numCols = numCols; + this.compressedData = compressedData; + this.compressionRatio = compressionRatio; + this.metadata = metadata; + } + + public CompressionType getType() { return type; } + public int getNumRows() { return numRows; } + public int getNumCols() { return numCols; } + public Object getCompressedData() { return compressedData; } + public double getCompressionRatio() { return compressionRatio; } + public byte[] getMetadata() { return metadata; } + + /** Estimate original size in bytes (8 bytes per double) */ + public long estimateOriginalSizeBytes() { + return (long) numRows * numCols * 8; + } + + /** Estimate compressed size in bytes */ + public long getCompressedSizeBytes() { + if (compressedData instanceof byte[]) { + return ((byte[]) compressedData).length; + } + return 0; + } + + @Override + public String toString() { + return String.format("CompressedMatrix[%s, %dx%d, ratio=%.2fx]", + type.getId(), numRows, numCols, compressionRatio); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionConfig.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionConfig.java new file mode 100644 index 00000000000..0f825570c16 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionConfig.java @@ -0,0 +1,110 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.controlprogram.federated.compression; + +import java.util.HashMap; +import java.util.Map; + +/** + * Immutable configuration for compression in federated operations. + * Uses the Builder pattern for flexible, readable configuration. + * + * Usage example: + * CompressionConfig config = CompressionConfig.builder() + * .enable(true) + * .withType(CompressionType.TOPK) + * .withSparsity(0.01) + * .build(); + */ +public class CompressionConfig { + + private final boolean enabled; + private final CompressionType type; + private final Map parameters; + + private CompressionConfig(Builder builder) { + this.enabled = builder.enabled; + this.type = builder.enabled ? builder.type : CompressionType.NONE; + this.parameters = new HashMap<>(builder.parameters); + } + + public boolean isEnabled() { return enabled; } + public CompressionType getType() { return type; } + public Map getParameters() { return new HashMap<>(parameters); } + + /** Convenience getter for sparsity parameter (TopK) */ + public double getSparsity() { + return (double) parameters.getOrDefault("sparsity", 0.01); + } + + /** Convenience getter for bits parameter (Quantization) */ + public int getBits() { + return (int) parameters.getOrDefault("bits", 4); + } + + @Override + public String toString() { + return String.format("CompressionConfig[enabled=%s, type=%s, params=%s]", + enabled, type.getId(), parameters); + } + + // ----------------------------------------------------------------------- + // Builder + // ----------------------------------------------------------------------- + + public static Builder builder() { + return new Builder(); + } + + public static class Builder { + private boolean enabled = false; + private CompressionType type = CompressionType.NONE; + private final Map parameters = new HashMap<>(); + + public Builder enable(boolean enabled) { + this.enabled = enabled; + return this; + } + + public Builder withType(CompressionType type) { + this.type = type; + return this; + } + + public Builder withParameter(String key, Object value) { + this.parameters.put(key, value); + return this; + } + + /** Shorthand for TopK sparsity ratio */ + public Builder withSparsity(double sparsity) { + return withParameter("sparsity", sparsity); + } + + /** Shorthand for quantization bit width */ + public Builder withBits(int bits) { + return withParameter("bits", bits); + } + + public CompressionConfig build() { + return new CompressionConfig(this); + } + } +} diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionFactory.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionFactory.java new file mode 100644 index 00000000000..0a6d1d0a3b2 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionFactory.java @@ -0,0 +1,94 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.controlprogram.federated.compression; + +import org.apache.sysds.runtime.controlprogram.federated.compression.TopK.TopKCompressor; +import org.apache.sysds.runtime.controlprogram.federated.compression.Quantization.ProbabilisticQuantizationCompressor; +import org.apache.sysds.runtime.controlprogram.federated.compression.exceptions.CompressionException; +import org.apache.sysds.runtime.controlprogram.federated.compression.exceptions.DecompressionException; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; + +/** + * Factory for creating compressor instances from configuration. + * Centralizes compressor instantiation and parameter validation. + * + * Usage: + * CompressionConfig config = CompressionConfig.builder() + * .enable(true) + * .withType(CompressionType.TOPK) + * .withSparsity(0.01) + * .build(); + * MatrixCompressor compressor = CompressionFactory.create(config); + */ +public class CompressionFactory { + + private CompressionFactory() { + // Utility class — no instantiation + } + + /** + * Create a compressor from a CompressionConfig. + * @param config The compression configuration + * @return A ready-to-use MatrixCompressor + * @throws IllegalArgumentException if the config is invalid + */ + public static MatrixCompressor create(CompressionConfig config) { + if(config == null || !config.isEnabled()) + return new PassthroughCompressor(); + switch(config.getType()) { + case TOPK: + return new TopKCompressor(config.getSparsity(), true); + case PROBABILISTIC_QUANTIZATION: + return new ProbabilisticQuantizationCompressor(config.getBits()); + case NONE: + default: + return new PassthroughCompressor(); + } + } + + // ----------------------------------------------------------------------- + // Passthrough compressor (no-op) for when compression is disabled + // ----------------------------------------------------------------------- + + /** + * No-op compressor: returns the matrix as-is. + * Used when compression is disabled or type is NONE. + */ + private static class PassthroughCompressor implements MatrixCompressor { + + @Override + public CompressedMatrix compress(MatrixBlock input) + throws CompressionException { + return new CompressedMatrix( + CompressionType.NONE, + input.getNumRows(), + input.getNumColumns(), + input, + 1.0 + ); + } + + @Override + public MatrixBlock decompress(CompressedMatrix compressed) + throws DecompressionException { + return (MatrixBlock) compressed.getCompressedData(); + } + } +} diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionType.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionType.java new file mode 100644 index 00000000000..ce8359a7f5a --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionType.java @@ -0,0 +1,57 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.controlprogram.federated.compression; + +/** + * Enumeration of supported compression techniques for federated learning. + * Used for configuration, serialization, and technique selection. + */ +public enum CompressionType { + + /** TopK sparsification: keep largest-magnitude elements only */ + TOPK("topk", "Top-K Sparsification"), + + /** Probabilistic quantization: reduce precision with stochastic rounding */ + PROBABILISTIC_QUANTIZATION("prob_quant", "Probabilistic Quantization"), + + /** No compression (passthrough) */ + NONE("none", "No Compression"); + + private final String id; + private final String description; + + CompressionType(String id, String description) { + this.id = id; + this.description = description; + } + + public String getId() { return id; } + public String getDescription() { return description; } + + /** Parse from string identifier (case-insensitive) */ + public static CompressionType fromString(String text) { + for (CompressionType type : CompressionType.values()) { + if (type.id.equalsIgnoreCase(text)) { + return type; + } + } + throw new IllegalArgumentException("Unknown compression type: " + text); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/MatrixCompressor.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/MatrixCompressor.java new file mode 100644 index 00000000000..301b245b0e3 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/MatrixCompressor.java @@ -0,0 +1,55 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.controlprogram.federated.compression; + +import org.apache.sysds.runtime.controlprogram.federated.compression.exceptions.CompressionException; +import org.apache.sysds.runtime.controlprogram.federated.compression.exceptions.DecompressionException; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; + +/** + * Interface for matrix compression techniques in federated learning. + * All compressors must implement compress/decompress operations. + */ +public interface MatrixCompressor { + + /** + * Compress a matrix block for transmission. + * @param input The source matrix to compress + * @return CompressedMatrix containing compressed data and metadata + * @throws CompressionException if compression fails + */ + CompressedMatrix compress(MatrixBlock input) throws CompressionException; + + /** + * Decompress a compressed matrix back to MatrixBlock. + * @param compressed The compressed data to decompress + * @return Reconstructed MatrixBlock (may be approximate) + * @throws DecompressionException if decompression fails + */ + MatrixBlock decompress(CompressedMatrix compressed) throws DecompressionException; + + /** + * Estimate the compression ratio achieved. + * Higher is better (e.g. 10.0 means 10x smaller). + */ + default double estimateCompressionRatio(long originalSize, long compressedSize) { + return compressedSize == 0 ? Double.MAX_VALUE : (double) originalSize / compressedSize; + } +} diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/Quantization/ProbabilisticQuantizationCompressor.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/Quantization/ProbabilisticQuantizationCompressor.java new file mode 100644 index 00000000000..7fcf9b9cf49 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/Quantization/ProbabilisticQuantizationCompressor.java @@ -0,0 +1,169 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.controlprogram.federated.compression.Quantization; + +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressedMatrix; +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressionType; +import org.apache.sysds.runtime.controlprogram.federated.compression.MatrixCompressor; +import org.apache.sysds.runtime.controlprogram.federated.compression.exceptions.CompressionException; +import org.apache.sysds.runtime.controlprogram.federated.compression.exceptions.DecompressionException; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; + +import java.util.Random; + +/** + * Probabilistic Quantization Compressor. + * + * Reduces numerical precision using stochastic rounding to maintain + * an unbiased estimator — meaning E[quantized] = original on average. + * This is critical for federated learning convergence guarantees. + * + * Supports 2, 4, or 8 bits per value: + * 2-bit → 4 levels → 16x compression vs 32-bit float + * 4-bit → 16 levels → 8x compression + * 8-bit → 256 levels → 4x compression + */ +public class ProbabilisticQuantizationCompressor implements MatrixCompressor { + + private final int bitsPerValue; // 2, 4, or 8 + private final Random rng; + + public ProbabilisticQuantizationCompressor(int bitsPerValue) { + if (bitsPerValue != 2 && bitsPerValue != 4 && bitsPerValue != 8) { + throw new IllegalArgumentException("bitsPerValue must be 2, 4, or 8"); + } + this.bitsPerValue = bitsPerValue; + this.rng = new Random(42); // Fixed seed for reproducibility + } + + @Override + public CompressedMatrix compress(MatrixBlock input) throws CompressionException { + try { + int numRows = input.getNumRows(); + int numCols = input.getNumColumns(); + int totalElements = numRows * numCols; + + // Find min and max for normalization + double[] minMax = findMinMax(input, numRows, numCols); + double min = minMax[0]; + double max = minMax[1]; + + int levels = 1 << bitsPerValue; // 2^bits + + // Quantize each element probabilistically + byte[] quantized = new byte[totalElements]; + for (int i = 0; i < numRows; i++) { + for (int j = 0; j < numCols; j++) { + double value = input.get(i, j); + quantized[i * numCols + j] = probabilisticRound(value, min, max, levels); + } + } + + double ratio = 32.0 / bitsPerValue; // vs 32-bit float + + QuantizedData data = new QuantizedData( + quantized, min, max, levels, bitsPerValue, numRows, numCols); + + return new CompressedMatrix( + CompressionType.PROBABILISTIC_QUANTIZATION, + numRows, numCols, data, ratio); + + } catch (Exception e) { + throw new CompressionException( + "Quantization compression failed: " + e.getMessage(), e); + } + } + + @Override + public MatrixBlock decompress(CompressedMatrix compressed) throws DecompressionException { + try { + QuantizedData data = (QuantizedData) compressed.getCompressedData(); + MatrixBlock result = new MatrixBlock(data.numRows, data.numCols, false); + result.allocateDenseBlock(); + + for (int i = 0; i < data.numRows; i++) { + for (int j = 0; j < data.numCols; j++) { + byte levelIndex = data.quantizedValues[i * data.numCols + j]; + double value = data.reconstructValue(levelIndex); + result.set(i, j, value); + } + } + + result.examSparsity(); + return result; + + } catch (ClassCastException e) { + throw new DecompressionException( + "Invalid compressed data type for Quantization", e); + } catch (Exception e) { + throw new DecompressionException( + "Quantization decompression failed: " + e.getMessage(), e); + } + } + + // ----------------------------------------------------------------------- + // Private helpers + // ----------------------------------------------------------------------- + + /** + * Stochastic rounding: for value x between levels q_i and q_{i+1}: + * P(round up) = (x - q_i) / (q_{i+1} - q_i) + * P(round down) = 1 - P(round up) + * This gives E[output] = x (unbiased). + */ + private byte probabilisticRound(double value, double min, double max, int levels) { + // Handle constant matrix edge case + if (max - min < 1e-10) return 0; + + // Normalize to [0, 1] + double normalized = (value - min) / (max - min); + + // Find bounding level indices + double scaled = normalized * (levels - 1); + int lowerIdx = (int) scaled; + int upperIdx = lowerIdx + 1; + + if (lowerIdx == upperIdx) { + return (byte) lowerIdx; + } + + // Probabilistic decision + double probUp = scaled - lowerIdx; + return (rng.nextDouble() < probUp) ? (byte) upperIdx : (byte) lowerIdx; + } + + /** Find min and max values across the entire matrix */ + private double[] findMinMax(MatrixBlock input, int numRows, int numCols) { + double min = Double.MAX_VALUE; + double max = -Double.MAX_VALUE; + + for (int i = 0; i < numRows; i++) { + for (int j = 0; j < numCols; j++) { + double val = input.get(i, j); + if (val < min) min = val; + if (val > max) max = val; + } + } + + // Handle all-zero matrix + if (min == Double.MAX_VALUE) { min = 0; max = 0; } + return new double[]{min, max}; + } +} diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/Quantization/QuantizedData.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/Quantization/QuantizedData.java new file mode 100644 index 00000000000..f1b2c2cbbe0 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/Quantization/QuantizedData.java @@ -0,0 +1,74 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.controlprogram.federated.compression.Quantization; + +import java.io.Serializable; + +/** + * Immutable container for probabilistically quantized matrix data. + * Stores quantized byte indices and the scaling parameters needed + * to reconstruct approximate original values on decompression. + */ +public class QuantizedData implements Serializable { + + private static final long serialVersionUID = 1L; + + public final byte[] quantizedValues; // Quantized level indices + public final double min; // Original minimum value + public final double max; // Original maximum value + public final int levels; // Number of quantization levels (2^bits) + public final int bitsPerValue; // Bits used per element + public final int numRows; + public final int numCols; + + public QuantizedData(byte[] quantizedValues, double min, double max, + int levels, int bitsPerValue, int numRows, int numCols) { + this.quantizedValues = quantizedValues.clone(); // Defensive copy + this.min = min; + this.max = max; + this.levels = levels; + this.bitsPerValue = bitsPerValue; + this.numRows = numRows; + this.numCols = numCols; + } + + /** Number of quantized elements */ + public int size() { + return quantizedValues.length; + } + + /** Estimate serialized size in bytes */ + public long estimateSizeBytes() { + return quantizedValues.length + 64; // +64 for scalar fields and headers + } + + /** Reconstruct a double value from a quantized level index */ + public double reconstructValue(byte levelIndex) { + if (max - min < 1e-10) return min; // Constant matrix + int idx = levelIndex & 0xFF; // Treat byte as unsigned + return min + (idx / (double)(levels - 1)) * (max - min); + } + + @Override + public String toString() { + return String.format("QuantizedData[%dx%d, levels=%d, bits=%d, min=%.4f, max=%.4f]", + numRows, numCols, levels, bitsPerValue, min, max); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/TopK/TopKCompressor.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/TopK/TopKCompressor.java new file mode 100644 index 00000000000..df3afc0b7f1 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/TopK/TopKCompressor.java @@ -0,0 +1,237 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.controlprogram.federated.compression.TopK; + +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressedMatrix; +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressionType; +import org.apache.sysds.runtime.controlprogram.federated.compression.MatrixCompressor; +import org.apache.sysds.runtime.controlprogram.federated.compression.exceptions.CompressionException; +import org.apache.sysds.runtime.controlprogram.federated.compression.exceptions.DecompressionException; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; + +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; +import java.util.PriorityQueue; + +/** + * TopK Sparsification Compressor. + * + * Keeps only the K largest-magnitude elements in the matrix, + * setting all others to zero. Optimal for gradient sparsification + * in federated learning where most gradient values are near zero. + * + * Compression ratio: approximately 1/sparsityRatio + * e.g. sparsityRatio=0.01 keeps 1% of elements → ~100x compression + */ +public class TopKCompressor implements MatrixCompressor { + + private final double sparsityRatio; // Fraction of elements to keep (0, 1] + private final boolean useHeap; // Use min-heap for O(n log k) selection + + /** + * @param sparsityRatio Fraction of elements to retain e.g. 0.01 = keep top 1% + * @param useHeap If true, use priority queue (faster for large matrices) + */ + public TopKCompressor(double sparsityRatio, boolean useHeap) { + if (sparsityRatio <= 0 || sparsityRatio > 1) { + throw new IllegalArgumentException("sparsityRatio must be in (0, 1]"); + } + this.sparsityRatio = sparsityRatio; + this.useHeap = useHeap; + } + + /** Default constructor: uses heap-based selection */ + public TopKCompressor(double sparsityRatio) { + this(sparsityRatio, true); + } + + @Override + public CompressedMatrix compress(MatrixBlock input) throws CompressionException { + try { + int numRows = input.getNumRows(); + int numCols = input.getNumColumns(); + int totalElements = numRows * numCols; + int k = (int) Math.max(1, Math.ceil(totalElements * sparsityRatio)); + + // If k covers everything, no compression needed + if (k >= totalElements) { + return new CompressedMatrix( + CompressionType.TOPK, numRows, numCols, input, 1.0); + } + + // Extract all non-zero elements with their linear indices + List elements = extractElements(input, numRows, numCols); + + // If fewer non-zeros than k, keep all of them + List topK = (elements.size() <= k) + ? new ArrayList<>(elements) + : selectTopK(elements, k); + + // Pack into TopKData + TopKData data = convertToTopKData(topK, numCols); + + double ratio = calculateCompressionRatio(totalElements, topK.size()); + + return new CompressedMatrix( + CompressionType.TOPK, numRows, numCols, data, ratio); + + } catch (Exception e) { + throw new CompressionException("TopK compression failed: " + e.getMessage(), e); + } + } + + @Override + public MatrixBlock decompress(CompressedMatrix compressed) throws DecompressionException { + try { + // Handle passthrough case (no compression was applied) + if (compressed.getCompressedData() instanceof MatrixBlock) { + return (MatrixBlock) compressed.getCompressedData(); + } + + TopKData data = (TopKData) compressed.getCompressedData(); + MatrixBlock result = new MatrixBlock( + compressed.getNumRows(), + compressed.getNumCols(), + true // Start sparse + ); + result.allocateSparseRowsBlock(); + + // Place values back at their original positions + for (int i = 0; i < data.indices.length; i++) { + int linearIdx = data.indices[i]; + int row = linearIdx / data.numCols; + int col = linearIdx % data.numCols; + result.set(row, col, data.values[i]); + } + + result.examSparsity(); + return result; + + } catch (ClassCastException e) { + throw new DecompressionException("Invalid compressed data type for TopK", e); + } catch (Exception e) { + throw new DecompressionException("TopK decompression failed: " + e.getMessage(), e); + } + } + + // ----------------------------------------------------------------------- + // Private helpers + // ----------------------------------------------------------------------- + + /** + * Extract all non-zero elements with their linear indices. + * Handles both dense and sparse MatrixBlock representations. + */ + private List extractElements(MatrixBlock input, int numRows, int numCols) { + List elements = new ArrayList<>(); + + if (input.isInSparseFormat()) { + // Sparse: iterate only over non-zero entries + for (int i = 0; i < numRows; i++) { + if (input.getSparseBlock() == null) continue; + if (input.getSparseBlock().isEmpty(i)) continue; + int[] rowIndices = input.getSparseBlock().indexes(i); + double[] rowValues = input.getSparseBlock().values(i); + int rowSize = input.getSparseBlock().size(i); + for (int j = 0; j < rowSize; j++) { + double val = rowValues[j]; + if (val != 0.0) { + int linearIdx = i * numCols + rowIndices[j]; + elements.add(new Element(linearIdx, val, Math.abs(val))); + } + } + } + } else { + // Dense: iterate all elements, skip zeros + double[] denseBlock = input.getDenseBlockValues(); + if (denseBlock != null) { + for (int i = 0; i < denseBlock.length; i++) { + if (denseBlock[i] != 0.0) { + elements.add(new Element(i, denseBlock[i], Math.abs(denseBlock[i]))); + } + } + } + } + return elements; + } + + /** + * Select top K elements by absolute value. + * Uses min-heap for O(n log k) when useHeap=true, + * or full sort O(n log n) otherwise. + */ + private List selectTopK(List elements, int k) { + if (useHeap) { + PriorityQueue minHeap = new PriorityQueue<>( + k, Comparator.comparingDouble(e -> e.absValue) + ); + for (Element e : elements) { + if (minHeap.size() < k) { + minHeap.offer(e); + } else if (e.absValue > minHeap.peek().absValue) { + minHeap.poll(); + minHeap.offer(e); + } + } + List result = new ArrayList<>(minHeap); + result.sort(Comparator.comparingInt(e -> e.index)); + return result; + } else { + elements.sort((a, b) -> Double.compare(b.absValue, a.absValue)); + return new ArrayList<>(elements.subList(0, k)); + } + } + + private TopKData convertToTopKData(List topK, int numCols) { + int[] indices = new int[topK.size()]; + double[] values = new double[topK.size()]; + for (int i = 0; i < topK.size(); i++) { + indices[i] = topK.get(i).index; + values[i] = topK.get(i).value; + } + return new TopKData(indices, values, numCols); + } + + private double calculateCompressionRatio(int total, int kept) { + if (kept == 0) return Double.MAX_VALUE; + // Original: total * 8 bytes (doubles) + // Compressed: kept * 12 bytes (int index + double value) + long originalBytes = (long) total * 8; + long compressedBytes = (long) kept * 12; + return (double) originalBytes / compressedBytes; + } + + // ----------------------------------------------------------------------- + // Inner class: element tracking during compression + // ----------------------------------------------------------------------- + + private static class Element { + final int index; // Linear index: row * numCols + col + final double value; // Original value + final double absValue; // Absolute value for magnitude comparison + + Element(int index, double value, double absValue) { + this.index = index; + this.value = value; + this.absValue = absValue; + } + } +} diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/TopK/TopKData.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/TopK/TopKData.java new file mode 100644 index 00000000000..5d7c6f7fa8f --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/TopK/TopKData.java @@ -0,0 +1,61 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.controlprogram.federated.compression.TopK; + +import java.io.Serializable; + +/** + * Immutable container for TopK-compressed matrix data. + * Stores only the K largest-magnitude elements with their positions, + * designed for efficient serialization across federated workers. + */ +public class TopKData implements Serializable { + + private static final long serialVersionUID = 1L; + + public final int[] indices; // Linear indices of kept elements (row*numCols + col) + public final double[] values; // Corresponding original values + public final int numCols; // Needed for index → (row, col) conversion + + public TopKData(int[] indices, double[] values, int numCols) { + if (indices.length != values.length) { + throw new IllegalArgumentException( + "Indices and values arrays must have the same length"); + } + this.indices = indices.clone(); // Defensive copy + this.values = values.clone(); + this.numCols = numCols; + } + + /** Number of kept elements */ + public int size() { + return indices.length; + } + + /** Estimate serialized size in bytes (4 bytes per int + 8 bytes per double) */ + public long estimateSizeBytes() { + return (long) indices.length * 12 + 64; // +64 for object headers + } + + @Override + public String toString() { + return String.format("TopKData[k=%d, numCols=%d]", indices.length, numCols); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/exceptions/CompressionException.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/exceptions/CompressionException.java new file mode 100644 index 00000000000..24574d7fb4d --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/exceptions/CompressionException.java @@ -0,0 +1,36 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.controlprogram.federated.compression.exceptions; + +/** + * Exception thrown when matrix compression fails. + */ +public class CompressionException extends Exception { + + private static final long serialVersionUID = 1L; + + public CompressionException(String message) { + super(message); + } + + public CompressionException(String message, Throwable cause) { + super(message, cause); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/exceptions/DecompressionException.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/exceptions/DecompressionException.java new file mode 100644 index 00000000000..7ba3601f107 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/exceptions/DecompressionException.java @@ -0,0 +1,36 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.controlprogram.federated.compression.exceptions; + +/** + * Exception thrown when matrix decompression fails. + */ +public class DecompressionException extends Exception { + + private static final long serialVersionUID = 1L; + + public DecompressionException(String message) { + super(message); + } + + public DecompressionException(String message, Throwable cause) { + super(message, cause); + } +} diff --git a/src/test/java/org/apache/sysds/test/functions/federated/io/ProbabilisticQuantizationCompressorTest.java b/src/test/java/org/apache/sysds/test/functions/federated/io/ProbabilisticQuantizationCompressorTest.java new file mode 100644 index 00000000000..20f9c433c2d --- /dev/null +++ b/src/test/java/org/apache/sysds/test/functions/federated/io/ProbabilisticQuantizationCompressorTest.java @@ -0,0 +1,142 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.sysds.test.functions.federated.io; + +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressedMatrix; +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressionType; +import org.apache.sysds.runtime.controlprogram.federated.compression.Quantization.ProbabilisticQuantizationCompressor; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; +import org.junit.Test; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +public class ProbabilisticQuantizationCompressorTest { + + @Test + public void testQuantizationBasicProperties() throws Exception { + MatrixBlock input = createRandomMatrix(10, 20); + ProbabilisticQuantizationCompressor compressor = + new ProbabilisticQuantizationCompressor(4); + CompressedMatrix compressed = compressor.compress(input); + MatrixBlock result = compressor.decompress(compressed); + + assertEquals(CompressionType.PROBABILISTIC_QUANTIZATION, compressed.getType()); + assertEquals(10, result.getNumRows()); + assertEquals(20, result.getNumColumns()); + assertEquals(8.0, compressed.getCompressionRatio(), 1e-10); + + double origMin = Double.MAX_VALUE; + double origMax = -Double.MAX_VALUE; + for(int i = 0; i < 10; i++) { + for(int j = 0; j < 20; j++) { + double v = input.get(i, j); + if(v < origMin) origMin = v; + if(v > origMax) origMax = v; + } + } + for(int i = 0; i < 10; i++) + for(int j = 0; j < 20; j++) { + double v = result.get(i, j); + assertTrue(v >= origMin - 1e-9); + assertTrue(v <= origMax + 1e-9); + } + } + + @Test + public void testFewerBitsGivesHigherRatio() throws Exception { + MatrixBlock input = createRandomMatrix(20, 20); + + double ratio2bit = new ProbabilisticQuantizationCompressor(2) + .compress(input).getCompressionRatio(); + double ratio8bit = new ProbabilisticQuantizationCompressor(8) + .compress(input).getCompressionRatio(); + + assertTrue(ratio2bit > ratio8bit); + } + + @Test + public void testUnbiasednessOverManyRuns() throws Exception { + MatrixBlock input = new MatrixBlock(1, 1, false); + input.allocateDenseBlock(); + input.set(0, 0, 5.0); + input.examSparsity(); + + double sum = 0.0; + int runs = 1000; + for(int r = 0; r < runs; r++) { + ProbabilisticQuantizationCompressor compressor = + new ProbabilisticQuantizationCompressor(4); + MatrixBlock result = compressor.decompress(compressor.compress(input)); + sum += result.get(0, 0); + } + assertEquals(5.0, sum / runs, 0.5); + } + + @Test + public void testConstantMatrix() throws Exception { + MatrixBlock input = new MatrixBlock(3, 3, false); + input.allocateDenseBlock(); + for(int i = 0; i < 3; i++) + for(int j = 0; j < 3; j++) + input.set(i, j, 7.0); + input.examSparsity(); + + ProbabilisticQuantizationCompressor compressor = + new ProbabilisticQuantizationCompressor(4); + MatrixBlock result = compressor.decompress(compressor.compress(input)); + + assertNotNull(result); + assertEquals(3, result.getNumRows()); + assertEquals(3, result.getNumColumns()); + } + + @Test(expected = IllegalArgumentException.class) + public void testInvalidBitsThrowsException() { + new ProbabilisticQuantizationCompressor(3); + } + + @Test + public void testCompressionRatio2Bit() throws Exception { + MatrixBlock input = createRandomMatrix(10, 10); + assertEquals(16.0, + new ProbabilisticQuantizationCompressor(2).compress(input).getCompressionRatio(), + 1e-10); + } + + @Test + public void testCompressionRatio8Bit() throws Exception { + MatrixBlock input = createRandomMatrix(10, 10); + assertEquals(4.0, + new ProbabilisticQuantizationCompressor(8).compress(input).getCompressionRatio(), + 1e-10); + } + + private MatrixBlock createRandomMatrix(int rows, int cols) { + MatrixBlock m = new MatrixBlock(rows, cols, false); + m.allocateDenseBlock(); + java.util.Random rng = new java.util.Random(42); + for(int i = 0; i < rows; i++) + for(int j = 0; j < cols; j++) + m.set(i, j, rng.nextGaussian() * 10); + m.examSparsity(); + return m; + } +} diff --git a/src/test/java/org/apache/sysds/test/functions/federated/io/TopKCompressorTest.java b/src/test/java/org/apache/sysds/test/functions/federated/io/TopKCompressorTest.java new file mode 100644 index 00000000000..3e07d802c05 --- /dev/null +++ b/src/test/java/org/apache/sysds/test/functions/federated/io/TopKCompressorTest.java @@ -0,0 +1,123 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.sysds.test.functions.federated.io; + +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressedMatrix; +import org.apache.sysds.runtime.controlprogram.federated.compression.CompressionType; +import org.apache.sysds.runtime.controlprogram.federated.compression.TopK.TopKCompressor; +import org.apache.sysds.runtime.matrix.data.MatrixBlock; +import org.junit.Test; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +public class TopKCompressorTest { + + @Test + public void testTopKBasicProperties() throws Exception { + MatrixBlock input = createRandomMatrix(10, 20); + TopKCompressor compressor = new TopKCompressor(0.1); + CompressedMatrix compressed = compressor.compress(input); + MatrixBlock result = compressor.decompress(compressed); + + assertEquals(CompressionType.TOPK, compressed.getType()); + assertEquals(10, result.getNumRows()); + assertEquals(20, result.getNumColumns()); + assertTrue(compressed.getCompressionRatio() > 0); + } + + @Test + public void testTopKKeepsLargestElements() throws Exception { + MatrixBlock input = new MatrixBlock(3, 3, false); + input.allocateDenseBlock(); + input.set(0, 0, 10.0); + input.set(1, 1, 5.0); + input.set(2, 2, 1.0); + input.examSparsity(); + + TopKCompressor compressor = new TopKCompressor(0.22); + CompressedMatrix compressed = compressor.compress(input); + MatrixBlock result = compressor.decompress(compressed); + + assertEquals(10.0, result.get(0, 0), 1e-10); + assertEquals(5.0, result.get(1, 1), 1e-10); + assertEquals(0.0, result.get(2, 2), 1e-10); + } + + @Test + public void testLowerSparsityGivesHigherRatio() throws Exception { + MatrixBlock input = createRandomMatrix(100, 100); + + TopKCompressor c1 = new TopKCompressor(0.1); + TopKCompressor c2 = new TopKCompressor(0.01); + + double ratio1 = c1.compress(input).getCompressionRatio(); + double ratio2 = c2.compress(input).getCompressionRatio(); + + assertTrue(ratio2 > ratio1); + } + + @Test + public void testAllZeroMatrix() throws Exception { + MatrixBlock input = new MatrixBlock(5, 5, false); + input.allocateDenseBlock(); + input.examSparsity(); + + TopKCompressor compressor = new TopKCompressor(0.1); + CompressedMatrix compressed = compressor.compress(input); + MatrixBlock result = compressor.decompress(compressed); + + for(int i = 0; i < 5; i++) + for(int j = 0; j < 5; j++) + assertEquals(0.0, result.get(i, j), 1e-10); + } + + @Test + public void testSparsityOfOneKeepsEverything() throws Exception { + MatrixBlock input = createRandomMatrix(5, 5); + TopKCompressor compressor = new TopKCompressor(1.0); + CompressedMatrix compressed = compressor.compress(input); + MatrixBlock result = compressor.decompress(compressed); + + for(int i = 0; i < 5; i++) + for(int j = 0; j < 5; j++) + assertEquals(input.get(i, j), result.get(i, j), 1e-10); + } + + @Test(expected = IllegalArgumentException.class) + public void testInvalidSparsityThrowsException() { + new TopKCompressor(0.0); + } + + @Test(expected = IllegalArgumentException.class) + public void testSparsityAboveOneThrowsException() { + new TopKCompressor(1.5); + } + + private MatrixBlock createRandomMatrix(int rows, int cols) { + MatrixBlock m = new MatrixBlock(rows, cols, false); + m.allocateDenseBlock(); + java.util.Random rng = new java.util.Random(42); + for(int i = 0; i < rows; i++) + for(int j = 0; j < cols; j++) + m.set(i, j, rng.nextGaussian() * 10); + m.examSparsity(); + return m; + } +}