diff --git a/src/main/java/dev/zarr/zarrjava/v3/codec/CodecBuilder.java b/src/main/java/dev/zarr/zarrjava/v3/codec/CodecBuilder.java index b599242..4e7c567 100644 --- a/src/main/java/dev/zarr/zarrjava/v3/codec/CodecBuilder.java +++ b/src/main/java/dev/zarr/zarrjava/v3/codec/CodecBuilder.java @@ -64,6 +64,16 @@ public CodecBuilder withTranspose(int[] order) { return this; } + public CodecBuilder withScaleOffset(Object offset, Object scale) { + codecs.add(new ScaleOffsetCodec(new ScaleOffsetCodec.Configuration(offset, scale))); + return this; + } + + public CodecBuilder withScaleOffset() { + codecs.add(new ScaleOffsetCodec(null)); + return this; + } + /** * Adds a {@code reshape} codec. Each entry of {@code shape} must be a positive {@link Integer}, the * special value {@code -1} (at most once), or an {@code int[]} / array of input dimension indices. diff --git a/src/main/java/dev/zarr/zarrjava/v3/codec/CodecRegistry.java b/src/main/java/dev/zarr/zarrjava/v3/codec/CodecRegistry.java index 0cc5544..505863a 100644 --- a/src/main/java/dev/zarr/zarrjava/v3/codec/CodecRegistry.java +++ b/src/main/java/dev/zarr/zarrjava/v3/codec/CodecRegistry.java @@ -19,6 +19,7 @@ public class CodecRegistry { addType("zstd", ZstdCodec.class); addType("crc32c", Crc32cCodec.class); addType("sharding_indexed", ShardingIndexedCodec.class); + addType("scale_offset", ScaleOffsetCodec.class); } public static void addType(String name, Class extends Codec> codecClass) { diff --git a/src/main/java/dev/zarr/zarrjava/v3/codec/core/ScaleOffsetCodec.java b/src/main/java/dev/zarr/zarrjava/v3/codec/core/ScaleOffsetCodec.java new file mode 100644 index 0000000..227e5b4 --- /dev/null +++ b/src/main/java/dev/zarr/zarrjava/v3/codec/core/ScaleOffsetCodec.java @@ -0,0 +1,379 @@ +package dev.zarr.zarrjava.v3.codec.core; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; +import dev.zarr.zarrjava.ZarrException; +import dev.zarr.zarrjava.core.ArrayMetadata.CoreArrayMetadata; +import dev.zarr.zarrjava.core.codec.ArrayArrayCodec; +import dev.zarr.zarrjava.v3.ArrayMetadata; +import dev.zarr.zarrjava.v3.DataType; +import dev.zarr.zarrjava.v3.codec.Codec; +import ucar.ma2.Array; +import ucar.ma2.IndexIterator; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import java.math.BigInteger; + +/** + * The {@code scale_offset} codec applies the affine transformation {@code (in - offset) * scale} to + * every array element on encode, and inverts it with {@code (in / scale) + offset} on decode. It is + * an {@code array -> array} codec: it does not change the data type or shape, it only rescales the + * stored values. It is typically followed by a narrowing codec (e.g. {@code cast_value}) that + * converts the rescaled array to a smaller data type to achieve (lossy) compression. + * + *
The arithmetic is performed using the semantics of the input array's data type. For integral + * data types the computation is exact: if any intermediate or final value is not representable in + * the data type (e.g. an unsigned subtraction going negative, an overflow, or a non-exact division + * on decode), the codec fails with a {@link ZarrException}. For floating-point data types the + * computation uses native {@code float}/{@code double} arithmetic, whose results are always + * representable (including {@code NaN} and {@code +-Infinity}). + * + *
The {@code offset} and {@code scale} configuration values are scalars encoded with the input + * array's data type using the Zarr V3 fill value encoding. A missing {@code offset} defaults to the + * additive identity (0); a missing {@code scale} defaults to the multiplicative identity (1). When + * both are absent the codec is a no-op. + * + *
Supported data types are the real-number types this library models: {@code int8/16/32/64}, + * {@code uint8/16/32/64}, {@code float32} and {@code float64}. Other data types from the codec + * specification (e.g. {@code float8_*}, {@code bfloat16}, {@code int2}) are not modelled here. + */ +public class ScaleOffsetCodec extends ArrayArrayCodec implements Codec { + + @JsonIgnore + @Nonnull + public final String name = "scale_offset"; + @Nullable + public final Configuration configuration; + + @JsonCreator + public ScaleOffsetCodec( + @Nullable @JsonProperty(value = "configuration") Configuration configuration + ) { + this.configuration = configuration; + } + + // ===== Codec pipeline integration ======================================================== + + @Override + public Array encode(Array chunkArray) throws ZarrException { + return transform(chunkArray, true); + } + + @Override + public Array decode(Array chunkArray) throws ZarrException { + return transform(chunkArray, false); + } + + @Override + public long computeEncodedSize(long inputByteLength, ArrayMetadata.CoreArrayMetadata arrayMetadata) + throws ZarrException { + // The data type and shape are unchanged, so the encoded chunk has the same byte length. + return inputByteLength; + } + + @Override + public CoreArrayMetadata resolveArrayMetadata() throws ZarrException { + super.resolveArrayMetadata(); + DataType type = arrayDataType(); + requireSupported(type); + // The data type stays the same; only the fill value is transformed (encode direction) so that + // fill-value-aware downstream codecs stay aligned with the rescaled data. + Object transformedFillValue = transformFillValue(arrayMetadata.parsedFillValue, type); + return new CoreArrayMetadata( + arrayMetadata.shape, arrayMetadata.chunkShape, type, transformedFillValue); + } + + private DataType arrayDataType() throws ZarrException { + if (!(arrayMetadata.dataType instanceof DataType)) { + throw new ZarrException("The scale_offset codec requires a Zarr v3 data type."); + } + return (DataType) arrayMetadata.dataType; + } + + // ===== Element transformation ============================================================ + + private Array transform(Array input, boolean encode) throws ZarrException { + DataType type = arrayDataType(); + requireSupported(type); + int[] shape = input.getShape(); + Array output = Array.factory(type.getMA2DataType(), shape); + IndexIterator in = input.getIndexIterator(); + IndexIterator out = output.getIndexIterator(); + + if (type == DataType.FLOAT32) { + float offset = floatParam(offsetConfig(), 0.0f); + float scale = floatParam(scaleConfig(), 1.0f); + while (in.hasNext()) { + float x = in.getFloatNext(); + out.setFloatNext(encode ? (x - offset) * scale : (x / scale) + offset); + } + } else if (type == DataType.FLOAT64) { + double offset = doubleParam(offsetConfig(), 0.0); + double scale = doubleParam(scaleConfig(), 1.0); + while (in.hasNext()) { + double x = in.getDoubleNext(); + out.setDoubleNext(encode ? (x - offset) * scale : (x / scale) + offset); + } + } else { + BigInteger offset = intParam(offsetConfig(), BigInteger.ZERO, type); + BigInteger scale = intParam(scaleConfig(), BigInteger.ONE, type); + BigInteger min = integerMin(type); + BigInteger max = integerMax(type); + while (in.hasNext()) { + BigInteger x = readInt(in, type); + BigInteger r = encode + ? encodeInt(x, offset, scale, type, min, max) + : decodeInt(x, offset, scale, type, min, max); + writeInt(out, type, r); + } + } + return output; + } + + private Object transformFillValue(Object fillValue, DataType type) throws ZarrException { + if (fillValue == null) { + return null; + } + if (type == DataType.FLOAT32) { + float offset = floatParam(offsetConfig(), 0.0f); + float scale = floatParam(scaleConfig(), 1.0f); + float x = ((Number) fillValue).floatValue(); + return (x - offset) * scale; + } + if (type == DataType.FLOAT64) { + double offset = doubleParam(offsetConfig(), 0.0); + double scale = doubleParam(scaleConfig(), 1.0); + double x = ((Number) fillValue).doubleValue(); + return (x - offset) * scale; + } + BigInteger offset = intParam(offsetConfig(), BigInteger.ZERO, type); + BigInteger scale = intParam(scaleConfig(), BigInteger.ONE, type); + BigInteger r = encodeInt(toBigInteger(fillValue, type), offset, scale, type, + integerMin(type), integerMax(type)); + return boxInt(r, type); + } + + // ===== Integer arithmetic (exact, with representability checks) ========================== + + private static BigInteger encodeInt(BigInteger x, BigInteger offset, BigInteger scale, + DataType type, BigInteger min, BigInteger max) + throws ZarrException { + BigInteger shifted = x.subtract(offset); + requireInRange(shifted, min, max, type, "intermediate value (in - offset)"); + BigInteger scaled = shifted.multiply(scale); + requireInRange(scaled, min, max, type, "result (in - offset) * scale"); + return scaled; + } + + private static BigInteger decodeInt(BigInteger x, BigInteger offset, BigInteger scale, + DataType type, BigInteger min, BigInteger max) + throws ZarrException { + if (scale.signum() == 0) { + throw new ZarrException("The scale_offset codec cannot decode with a scale of 0."); + } + BigInteger[] quotientRemainder = x.divideAndRemainder(scale); + if (quotientRemainder[1].signum() != 0) { + throw new ZarrException( + "The scale_offset codec cannot decode the value " + x + " because it is not exactly " + + "divisible by the scale " + scale + " in the '" + type.getValue() + "' data type."); + } + BigInteger divided = quotientRemainder[0]; + requireInRange(divided, min, max, type, "intermediate value (in / scale)"); + BigInteger result = divided.add(offset); + requireInRange(result, min, max, type, "result (in / scale) + offset"); + return result; + } + + private static void requireInRange(BigInteger value, BigInteger min, BigInteger max, + DataType type, String label) throws ZarrException { + if (value.compareTo(min) < 0 || value.compareTo(max) > 0) { + throw new ZarrException( + "The scale_offset " + label + " (" + value + ") is not representable in the '" + + type.getValue() + "' data type."); + } + } + + // ===== Configuration parameter parsing =================================================== + + @Nullable + private Object offsetConfig() { + return configuration == null ? null : configuration.offset; + } + + @Nullable + private Object scaleConfig() { + return configuration == null ? null : configuration.scale; + } + + private static float floatParam(@Nullable Object raw, float identity) throws ZarrException { + if (raw == null) { + return identity; + } + return ((Number) ArrayMetadata.parseFillValue(raw, DataType.FLOAT32)).floatValue(); + } + + private static double doubleParam(@Nullable Object raw, double identity) throws ZarrException { + if (raw == null) { + return identity; + } + return ((Number) ArrayMetadata.parseFillValue(raw, DataType.FLOAT64)).doubleValue(); + } + + private static BigInteger intParam(@Nullable Object raw, BigInteger identity, DataType type) + throws ZarrException { + if (raw == null) { + return identity; + } + return toBigInteger(ArrayMetadata.parseFillValue(raw, type), type); + } + + // ===== Data type facts and element reading/writing ======================================= + + private static void requireSupported(DataType type) throws ZarrException { + if (type == DataType.BOOL) { + throw new ZarrException( + "The scale_offset codec does not support the data type '" + type.getValue() + + "'. Supported types are the integral and floating-point real-number types."); + } + } + + private static BigInteger readInt(IndexIterator it, DataType type) { + switch (type) { + case INT8: + return BigInteger.valueOf(it.getByteNext()); + case UINT8: + return BigInteger.valueOf(it.getByteNext() & 0xFFL); + case INT16: + return BigInteger.valueOf(it.getShortNext()); + case UINT16: + return BigInteger.valueOf(it.getShortNext() & 0xFFFFL); + case INT32: + return BigInteger.valueOf(it.getIntNext()); + case UINT32: + return BigInteger.valueOf(it.getIntNext() & 0xFFFFFFFFL); + case INT64: + return BigInteger.valueOf(it.getLongNext()); + case UINT64: + return new BigInteger(Long.toUnsignedString(it.getLongNext())); + default: + throw new IllegalStateException("Unsupported scale_offset data type: " + type); + } + } + + private static void writeInt(IndexIterator it, DataType type, BigInteger value) { + switch (type) { + case INT8: + case UINT8: + it.setByteNext(value.byteValue()); + break; + case INT16: + case UINT16: + it.setShortNext(value.shortValue()); + break; + case INT32: + case UINT32: + it.setIntNext(value.intValue()); + break; + case INT64: + case UINT64: + it.setLongNext(value.longValue()); + break; + default: + throw new IllegalStateException("Unsupported scale_offset data type: " + type); + } + } + + private static BigInteger toBigInteger(Object boxed, DataType type) { + Number number = (Number) boxed; + switch (type) { + case INT8: + return BigInteger.valueOf(number.byteValue()); + case UINT8: + return BigInteger.valueOf(number.longValue() & 0xFFL); + case INT16: + return BigInteger.valueOf(number.shortValue()); + case UINT16: + return BigInteger.valueOf(number.longValue() & 0xFFFFL); + case INT32: + return BigInteger.valueOf(number.intValue()); + case UINT32: + return BigInteger.valueOf(number.longValue() & 0xFFFFFFFFL); + case INT64: + return BigInteger.valueOf(number.longValue()); + case UINT64: + return new BigInteger(Long.toUnsignedString(number.longValue())); + default: + throw new IllegalStateException("Unsupported scale_offset data type: " + type); + } + } + + private static Object boxInt(BigInteger value, DataType type) { + switch (type) { + case INT8: + case UINT8: + return value.byteValue(); + case INT16: + case UINT16: + return value.shortValue(); + case INT32: + case UINT32: + return value.intValue(); + case INT64: + case UINT64: + return value.longValue(); + default: + throw new IllegalStateException("Unsupported scale_offset data type: " + type); + } + } + + private static int integerBits(DataType type) { + return type.getByteCount() * 8; + } + + private static boolean isUnsigned(DataType type) { + return type == DataType.UINT8 || type == DataType.UINT16 || type == DataType.UINT32 + || type == DataType.UINT64; + } + + private static BigInteger integerMin(DataType type) { + if (isUnsigned(type)) { + return BigInteger.ZERO; + } + return BigInteger.ONE.shiftLeft(integerBits(type) - 1).negate(); + } + + private static BigInteger integerMax(DataType type) { + if (isUnsigned(type)) { + return BigInteger.ONE.shiftLeft(integerBits(type)).subtract(BigInteger.ONE); + } + return BigInteger.ONE.shiftLeft(integerBits(type) - 1).subtract(BigInteger.ONE); + } + + // ===== Configuration ===================================================================== + + public static final class Configuration { + + /** The offset subtracted on encode, as a JSON scalar in the input array's data type. */ + @Nullable + @JsonInclude(JsonInclude.Include.NON_NULL) + @JsonProperty("offset") + public final Object offset; + + /** The scale multiplied on encode, as a JSON scalar in the input array's data type. */ + @Nullable + @JsonInclude(JsonInclude.Include.NON_NULL) + @JsonProperty("scale") + public final Object scale; + + @JsonCreator + public Configuration( + @Nullable @JsonProperty("offset") Object offset, + @Nullable @JsonProperty("scale") Object scale) { + this.offset = offset; + this.scale = scale; + } + } +} diff --git a/src/test/java/dev/zarr/zarrjava/codec/ScaleOffsetCodecTest.java b/src/test/java/dev/zarr/zarrjava/codec/ScaleOffsetCodecTest.java new file mode 100644 index 0000000..359df11 --- /dev/null +++ b/src/test/java/dev/zarr/zarrjava/codec/ScaleOffsetCodecTest.java @@ -0,0 +1,123 @@ +package dev.zarr.zarrjava.codec; + +import dev.zarr.zarrjava.ZarrException; +import dev.zarr.zarrjava.ZarrTest; +import dev.zarr.zarrjava.store.FilesystemStore; +import dev.zarr.zarrjava.store.StoreHandle; +import dev.zarr.zarrjava.v3.Array; +import dev.zarr.zarrjava.v3.ArrayMetadata; +import dev.zarr.zarrjava.v3.ArrayMetadataBuilder; +import dev.zarr.zarrjava.v3.DataType; +import dev.zarr.zarrjava.v3.codec.core.ScaleOffsetCodec; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import ucar.ma2.MAMath; + +import java.io.IOException; + +import static dev.zarr.zarrjava.core.ArrayMetadata.parseFillValue; +import static dev.zarr.zarrjava.utils.Utils.toLongArray; +import static org.junit.Assert.assertThrows; + +public class ScaleOffsetCodecTest extends ZarrTest { + + private static ScaleOffsetCodec scaleOffsetCodec(Object offset, Object scale, DataType dataType, + Object fillValue, int[] shape) throws ZarrException { + ScaleOffsetCodec codec = new ScaleOffsetCodec(new ScaleOffsetCodec.Configuration(offset, scale)); + codec.setCoreArrayMetadata(new ArrayMetadata.CoreArrayMetadata( + toLongArray(shape), shape, dataType, parseFillValue(fillValue, dataType))); + return codec; + } + + @Test + public void testScaleOffsetCodecFloat() throws ZarrException { + // scale 0.5 is exactly representable in float32, so the round-trip is lossless here. + ucar.ma2.Array in = ucar.ma2.Array.factory(ucar.ma2.DataType.FLOAT, new int[]{4}, + new float[]{4.0f, 5.0f, 6.0f, 8.0f}); + // (x - 5) * 0.5 + ucar.ma2.Array encoded = ucar.ma2.Array.factory(ucar.ma2.DataType.FLOAT, new int[]{4}, + new float[]{-0.5f, 0.0f, 0.5f, 1.5f}); + + ScaleOffsetCodec codec = scaleOffsetCodec(5, 0.5, DataType.FLOAT32, null, new int[]{4}); + assert MAMath.equals(encoded, codec.encode(in.copy())); + assert MAMath.equals(in, codec.decode(encoded.copy())); + } + + @Test + public void testScaleOffsetCodecUintOffsetOnly() throws ZarrException { + // Range reduction: subtract 1000, leaving values that fit in a byte. scale defaults to 1. + ucar.ma2.Array in = ucar.ma2.Array.factory(ucar.ma2.DataType.USHORT, new int[]{4}, + new short[]{1000, 1001, 1050, 1255}); + ucar.ma2.Array encoded = ucar.ma2.Array.factory(ucar.ma2.DataType.USHORT, new int[]{4}, + new short[]{0, 1, 50, 255}); + + ScaleOffsetCodec codec = scaleOffsetCodec(1000, null, DataType.UINT16, null, new int[]{4}); + assert MAMath.equals(encoded, codec.encode(in.copy())); + assert MAMath.equals(in, codec.decode(encoded.copy())); + } + + @Test + public void testScaleOffsetCodecNoOp() throws ZarrException { + ucar.ma2.Array in = ucar.ma2.Array.factory(ucar.ma2.DataType.INT, new int[]{3}, + new int[]{-7, 0, 42}); + ScaleOffsetCodec codec = new ScaleOffsetCodec(null); + codec.setCoreArrayMetadata(new ArrayMetadata.CoreArrayMetadata( + new long[]{3}, new int[]{3}, DataType.INT32, null)); + assert MAMath.equals(in, codec.encode(in.copy())); + assert MAMath.equals(in, codec.decode(in.copy())); + } + + @Test + public void testScaleOffsetCodecFillValueTransform() throws ZarrException { + // The fill value is transformed with the encode formula and reported downstream. + ScaleOffsetCodec codec = scaleOffsetCodec(5, 0.5, DataType.FLOAT32, 5.0f, new int[]{4}); + Object resolvedFill = codec.resolveArrayMetadata().parsedFillValue; + Assertions.assertEquals(0.0f, resolvedFill); + + ScaleOffsetCodec uintCodec = scaleOffsetCodec(1000, null, DataType.UINT16, 1000, new int[]{4}); + Assertions.assertEquals((short) 0, uintCodec.resolveArrayMetadata().parsedFillValue); + } + + @Test + public void testScaleOffsetCodecIntegerOutOfRangeIsError() throws ZarrException { + // 500 - 1000 = -500 is not representable in uint16 -> hard error (no numpy-style wraparound). + ucar.ma2.Array in = ucar.ma2.Array.factory(ucar.ma2.DataType.USHORT, new int[]{1}, + new short[]{500}); + ScaleOffsetCodec codec = scaleOffsetCodec(1000, null, DataType.UINT16, null, new int[]{1}); + assertThrows(ZarrException.class, () -> codec.encode(in)); + } + + @Test + public void testScaleOffsetCodecNonExactDivisionIsError() throws ZarrException { + // Decoding requires in / scale to be an exact integer for integral data types. + ucar.ma2.Array stored = ucar.ma2.Array.factory(ucar.ma2.DataType.INT, new int[]{1}, + new int[]{5}); + ScaleOffsetCodec codec = scaleOffsetCodec(0, 10, DataType.INT32, null, new int[]{1}); + assertThrows(ZarrException.class, () -> codec.decode(stored)); + } + + @Test + public void testScaleOffsetCodecReadWrite() throws IOException, ZarrException { + StoreHandle storeHandle = new FilesystemStore(TESTOUTPUT).resolve("testScaleOffsetCodecReadWrite"); + float[] values = new float[16 * 16 * 16]; + for (int i = 0; i < values.length; i++) { + // multiples of 0.5, all exactly representable and exactly recoverable with scale 0.5 + values[i] = (i % 32) * 0.5f + 3.0f; + } + ucar.ma2.Array testData = ucar.ma2.Array.factory(ucar.ma2.DataType.FLOAT, + new int[]{16, 16, 16}, values); + + ArrayMetadataBuilder builder = Array.metadataBuilder() + .withShape(16, 16, 16) + .withDataType(DataType.FLOAT32) + .withChunkShape(4, 8, 16) + .withFillValue(3.0f) + .withCodecs(c -> c.withScaleOffset(3.0, 0.5)); + Array writeArray = Array.create(storeHandle, builder.build()); + writeArray.write(testData); + + Array readArray = Array.open(storeHandle); + ucar.ma2.Array result = readArray.read(); + assert MAMath.equals(testData, result); + } +}