[feat] enable shuffle for all codecs,accept info in zmat.zmat

Author: fangq

python/zmat/__init__.py

@@ -16,6 +16,9 @@
 NumPy-aware API:
     compressed, info = zmat.compress(arr, info=True)
     restored_arr     = zmat.decompress(compressed, info=info)
+
+    compressed, info = zmat.zmat(arr, info=True)          # low-level with info
+    restored_arr     = zmat.zmat(compressed, info=info)   # low-level restore
 """
 
 from _zmat import compress as _compress
@@ -26,9 +29,29 @@
 
 __all__ = ["compress", "decompress", "encode", "decode", "zmat"]
 
+
+def _byte_shuffle(data_bytes, typesize):
+    """Regroup bytes by position within each element (byte-shuffle filter).
+
+    Transforms [b0_e1 b1_e1 ... b0_e2 b1_e2 ...] into
+    [b0_e1 b0_e2 ... b1_e1 b1_e2 ...], improving compressibility of
+    structured numerical arrays with any codec.
+    """
+    import numpy as np
+    arr = np.frombuffer(data_bytes, dtype=np.uint8).reshape(-1, typesize)
+    return arr.flatten(order='F').tobytes()
+
+
+def _byte_unshuffle(data_bytes, typesize):
+    """Reverse of _byte_shuffle."""
+    import numpy as np
+    nelems = len(data_bytes) // typesize
+    arr = np.frombuffer(data_bytes, dtype=np.uint8).reshape(typesize, nelems)
+    return arr.flatten(order='F').tobytes()
+
 __version__ = "1.0.2"
 
-def compress(data, method="zlib", level=1, info=False):
+def compress(data, method="zlib", level=1, info=False, shuffle=0):
     """Compress *data* using the requested algorithm.
 
     Parameters
@@ -57,10 +80,19 @@ def compress(data, method="zlib", level=1, info=False):
         - ``'byte'``   — bytes per element (``data.itemsize``)
         - ``'method'`` — the compression method used
         - ``'order'``  — ``'F'`` for Fortran-contiguous, ``'C'`` otherwise
+        - ``'shuffle'``— byte-shuffle level applied (0 = none, 1 = byte)
+        - ``'typesize'``— element size in bytes used for shuffle
 
         When *info=True* but *data* is not an ndarray, the tuple
         ``(compressed_bytes, None)`` is returned so callers can always
         unpack two values.
+    shuffle : int, optional
+        Byte-shuffle level for non-blosc2 codecs (default ``0`` = disabled,
+        ``1`` = byte-shuffle).  Requires *info=True* so that the shuffle
+        state can be recorded and reversed on decompression.  For blosc2
+        methods the shuffle is handled by the C layer and this parameter
+        is ignored.  Has no effect when *data* is not a
+        :class:`numpy.ndarray`.
 
     Returns
     -------
@@ -76,29 +108,36 @@ def compress(data, method="zlib", level=1, info=False):
         compressed = zmat.compress(b"hello " * 1000)
         original   = zmat.decompress(compressed)
 
-    NumPy array round-trip::
+    NumPy array round-trip with byte-shuffle::
 
         import numpy as np
         arr = np.random.rand(100, 100)
-        compressed, info = zmat.compress(arr, info=True)
-        restored = zmat.decompress(compressed, info=info)   # ndarray
+        compressed, info = zmat.compress(arr, method='lz4', info=True, shuffle=1)
+        restored = zmat.decompress(compressed, info=info)
         assert np.array_equal(restored, arr)
     """
+    _use_shuffle = (shuffle > 0 and "blosc2" not in method and method != "base64")
+
     if info:
         try:
             import numpy as np
 
             if isinstance(data, np.ndarray):
                 order = "F" if np.isfortran(data) else "C"
+                ts = data.itemsize
+                apply_shuffle = _use_shuffle and ts > 1
                 arr_info = {
                     "type": str(data.dtype),
                     "shape": tuple(data.shape),
-                    "byte": data.itemsize,
+                    "byte": ts,
                     "method": method,
                     "order": order,
+                    "shuffle": shuffle if apply_shuffle else 0,
+                    "typesize": ts,
                 }
-                # always serialise as C-contiguous bytes
                 flat = np.ascontiguousarray(data).tobytes()
+                if apply_shuffle:
+                    flat = _byte_shuffle(flat, ts)
                 compressed = _compress(flat, method=method, level=level)
                 return compressed, arr_info
         except ImportError:
@@ -150,6 +189,13 @@ def decompress(data, method="zlib", info=None):
     if info is not None:
         actual_method = info.get("method", method)
         raw = _decompress(data, method=actual_method)
+
+        # unshuffle if compression applied wrapper-level byte shuffle
+        shuf = info.get("shuffle", 0)
+        ts   = info.get("typesize", 1)
+        if shuf > 0 and ts > 1 and "blosc2" not in actual_method:
+            raw = _byte_unshuffle(raw, ts)
+
         try:
             import numpy as np
 
@@ -167,17 +213,21 @@ def decompress(data, method="zlib", info=None):
     return _decompress(data, method=method)
 
 
-def zmat(data, iscompress=1, method="zlib", nthread=1, shuffle=1, typesize=4):
+def zmat(data, iscompress=1, method="zlib", nthread=1, shuffle=1, typesize=4, info=False):
     """Low-level compression/decompression interface with full parameter control.
 
+    Mirrors the MATLAB ``[ss, info] = zmat(arr)`` / ``zmat(ss, info)`` pattern
+    when *info* is used.
+
     Parameters
     ----------
-    data : bytes, bytearray, or buffer-protocol object
+    data : bytes, bytearray, buffer-protocol object, or numpy.ndarray
         Input data.
     iscompress : int
         ``1`` to compress at the default level (default), ``0`` to
         decompress, or a negative integer to set the compression level
-        (e.g. ``-9`` for maximum compression).
+        (e.g. ``-9`` for maximum compression).  Ignored when *info* is a
+        dict — decompression is performed automatically in that case.
     method : str
         Compression algorithm (default ``'zlib'``).
     nthread : int
@@ -188,12 +238,105 @@ def zmat(data, iscompress=1, method="zlib", nthread=1, shuffle=1, typesize=4):
     typesize : int
         Element byte size used by the blosc2 byte-shuffle filter
         (default ``4``).
+    info : bool or dict, optional
+        * ``False`` (default) — plain bytes in, bytes out.
+        * ``True`` — when *data* is a :class:`numpy.ndarray`, capture its
+          dtype, shape, and memory order and return ``(compressed_bytes,
+          info_dict)``.  When *data* is not an ndarray, returns
+          ``(compressed_bytes, None)``.
+        * ``dict`` — treat as the info dict previously returned by this
+          function or by :func:`compress`.  Decompresses *data* and
+          reconstructs the original :class:`numpy.ndarray` using the
+          stored metadata.  The method is taken from ``info['method']``;
+          the *method* argument is used only as a fallback.
 
     Returns
     -------
     bytes
-        Compressed or decompressed data.
+        Compressed or decompressed data when *info* is ``False``.
+    tuple[bytes, dict | None]
+        ``(compressed_bytes, info_dict)`` when *info=True*.
+    numpy.ndarray
+        Reconstructed array when *info* is a dict and NumPy is available.
+
+    Examples
+    --------
+    NumPy array round-trip::
+
+        import numpy as np
+        arr = np.random.rand(100, 100).astype(np.float32)
+        compressed, info = zmat.zmat(arr, info=True)
+        restored = zmat.zmat(compressed, info=info)
+        assert np.array_equal(restored, arr)
+
+    blosc2 with multi-threading::
+
+        out = zmat.zmat(data, iscompress=1, method='blosc2zstd',
+                        nthread=4, shuffle=1, typesize=8)
     """
+    _use_shuffle = (shuffle > 0 and "blosc2" not in method and method != "base64")
+
+    # info dict supplied → decompress and reconstruct numpy array
+    if isinstance(info, dict):
+        actual_method = info.get("method", method)
+        # blosc2 shuffle is handled by the C layer; pass it through unchanged
+        c_shuffle = shuffle if "blosc2" in actual_method else 0
+        raw = _zmat_c(data, iscompress=0, method=actual_method,
+                      nthread=nthread, shuffle=c_shuffle, typesize=typesize)
+
+        # unshuffle if wrapper-level shuffle was recorded in info
+        shuf = info.get("shuffle", 0)
+        ts   = info.get("typesize", 1)
+        if shuf > 0 and ts > 1 and "blosc2" not in actual_method:
+            raw = _byte_unshuffle(raw, ts)
+
+        try:
+            import numpy as np
+
+            dtype = np.dtype(info["type"])
+            shape = tuple(info["shape"])
+            order = info.get("order", "C")
+            arr = np.frombuffer(raw, dtype=dtype).copy().reshape(shape)
+            if order == "F":
+                arr = np.asfortranarray(arr)
+            return arr
+        except ImportError:
+            return raw
+
+    # info=True → capture numpy array metadata before compressing
+    if info is True:
+        try:
+            import numpy as np
+
+            if isinstance(data, np.ndarray):
+                order = "F" if np.isfortran(data) else "C"
+                ts = data.itemsize
+                apply_shuffle = _use_shuffle and ts > 1
+                arr_info = {
+                    "type": str(data.dtype),
+                    "shape": tuple(data.shape),
+                    "byte": ts,
+                    "method": method,
+                    "order": order,
+                    "shuffle": shuffle if apply_shuffle else 0,
+                    "typesize": ts,
+                }
+                flat = np.ascontiguousarray(data).tobytes()
+                if apply_shuffle:
+                    flat = _byte_shuffle(flat, ts)
+                # for blosc2, pass shuffle/typesize to C; for others, already done
+                c_shuffle  = shuffle if "blosc2" in method else 0
+                c_typesize = typesize if "blosc2" in method else 1
+                compressed = _zmat_c(flat, iscompress=iscompress, method=method,
+                                     nthread=nthread, shuffle=c_shuffle, typesize=c_typesize)
+                return compressed, arr_info
+        except ImportError:
+            pass
+
+        # non-ndarray with info=True: compress normally, return (bytes, None)
+        return _zmat_c(data, iscompress=iscompress, method=method,
+                       nthread=nthread, shuffle=shuffle, typesize=typesize), None
+
     return _zmat_c(
         data,
         iscompress=iscompress,

zmat.m

@@ -46,7 +46,12 @@
 %     options: a series of ('name', value) pairs, supported options include
 %             'nthread': followed by an integer specifying number of threads for blosc2 meta-compressors
 %             'typesize': followed by an integer specifying the number of bytes per data element (used for shuffle)
-%             'shuffle': shuffle methods in blosc2 meta-compressor, 0 disable, 1, byte-shuffle
+%             'shuffle': 0 to disable (default for non-blosc2), 1 to enable byte-shuffle.
+%                     For blosc2 methods the shuffle is applied inside the C layer.
+%                     For all other methods, when [output,info]=zmat(...) is called,
+%                     the shuffle is applied in MATLAB before compression and reversed
+%                     after decompression; the info struct records 'shuffle' and
+%                     'typesize' so that zmat(compressed, info) restores the original array.
 %
 % output:
 %      output: a uint8 row vector, storing the compressed or decompressed data;
@@ -220,7 +225,34 @@
     iscompress = -9;
 end
 
-[varargout{1:max(1, nargout)}] = zipmat(input, iscompress, zipmethod, nthread, shuffle, typesize);
+%% wrapper-level byte shuffle: applies to non-blosc2 codecs when info is requested
+do_wrapper_shuffle = (nargout > 1 && shuffle > 0 && iscompress ~= 0 && ...
+                      isempty(strfind(zipmethod, 'blosc2')) && ...
+                      ~strcmp(zipmethod, 'base64') && ...
+                      isempty(specialtype) && typesize > 1);
+
+if (do_wrapper_shuffle)
+    %% convert input to bytes, apply byte shuffle, then compress
+    orig_class = class(input);
+    orig_size  = size(input);
+    if (~isa(input, 'uint8'))
+        raw_bytes = typecast(input(:)', 'uint8');
+    else
+        raw_bytes = input(:)';
+    end
+    nelems = numel(raw_bytes) / typesize;
+    M = reshape(raw_bytes, typesize, nelems);   % typesize x nelems: col = one element
+    shuffled_bytes = reshape(M', 1, []);        % flatten row-major: all byte-0s, then byte-1s...
+    [varargout{1:max(1, nargout)}] = zipmat(shuffled_bytes, iscompress, zipmethod, nthread, 0, 1);
+    %% overwrite info with original array metadata and record shuffle state
+    varargout{2}.type     = orig_class;
+    varargout{2}.size     = orig_size;
+    varargout{2}.byte     = typesize;
+    varargout{2}.shuffle  = shuffle;
+    varargout{2}.typesize = typesize;
+else
+    [varargout{1:max(1, nargout)}] = zipmat(input, iscompress, zipmethod, nthread, shuffle, typesize);
+end
 
 %% store special matrix type info in the output info struct
 if (nargout > 1 && ~isempty(specialtype))
@@ -295,6 +327,17 @@
         end
     end
 
+    %% unshuffle if compression applied wrapper-level byte shuffle
+    if (isfield(inputinfo, 'shuffle') && inputinfo.shuffle > 0 && ...
+        isfield(inputinfo, 'typesize') && inputinfo.typesize > 1 && ...
+        isempty(strfind(inputinfo.method, 'blosc2')))
+        ts = inputinfo.typesize;
+        raw_bytes = varargout{1}(:)';
+        nelems = numel(raw_bytes) / ts;
+        M = reshape(raw_bytes, nelems, ts);     % nelems x ts: col = one byte position
+        varargout{1} = reshape(M', 1, []);      % flatten row-major: bytes of e1, then e2...
+    end
+
     if (strcmp(inputinfo.type, 'logical'))
         varargout{1} = logical(varargout{1});
     else