biotite 1.1.0__cp313-cp313-win_amd64.whl

biotite/structure/io/pdbx/component.py

@@ -0,0 +1,245 @@
+# This source code is part of the Biotite package and is distributed
+# under the 3-Clause BSD License. Please see 'LICENSE.rst' for further
+# information.
+
+"""
+This module contains internally abstract classes for representing parts
+of CIF/BinaryCIF files, such as categories and columns.
+"""
+
+__name__ = "biotite.structure.io.pdbx"
+__author__ = "Patrick Kunzmann"
+__all__ = ["MaskValue"]
+
+from abc import ABCMeta, abstractmethod
+from collections.abc import MutableMapping
+from enum import IntEnum
+from biotite.file import DeserializationError, SerializationError
+
+
+class MaskValue(IntEnum):
+    """
+    This enum type represents the possible values of a mask array.
+
+    - `PRESENT` : A value is present.
+    - `INAPPLICABLE` : For this row no value is applicable or
+      inappropriate (``.`` in *CIF*).
+      In some cases it may also refer to a default value for the
+      respective column.
+    - `MISSING` : For this row the value is missing or unknown
+      (``?`` in *CIF*).
+    """
+
+    PRESENT = 0
+    INAPPLICABLE = 1
+    MISSING = 2
+
+
+class _Component(metaclass=ABCMeta):
+    """
+    Base class for all components in a CIF/BinaryCIF file.
+    """
+
+    @staticmethod
+    def subcomponent_class():
+        """
+        Get the class of the components that are stored in this component.
+
+        Returns
+        -------
+        subcomponent_class : type
+            The class of the subcomponent.
+            If this component already represents the lowest level, i.e.
+            it does not contain subcomponents, ``None`` is
+            returned.
+        """
+        return None
+
+    @staticmethod
+    def supercomponent_class():
+        """
+        Get the class of the component that contains this component.
+
+        Returns
+        -------
+        supercomponent_class : type
+            The class of the supercomponent.
+            If this component present already the highest level, i.e.
+            it is not contained in another component, ``None`` is
+            returned.
+        """
+        return None
+
+    @staticmethod
+    @abstractmethod
+    def deserialize(content):
+        """
+        Create this component by deserializing the given content.
+
+        Parameters
+        ----------
+        content : str or dict
+            The content to be deserialized.
+            The type of this parameter depends on the file format.
+            In case of *CIF* files, this is the text of the lines
+            that represent this component.
+            In case of *BinaryCIF* files, this is a dictionary
+            parsed from the *MessagePack* data.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def serialize(self):
+        """
+        Convert this component into a Python object that can be written
+        to a file.
+
+        Returns
+        -------
+        content : str or dict
+            The content to be serialized.
+            The type of this return value depends on the file format.
+            In case of *CIF* files, this is the text of the lines
+            that represent this component.
+            In case of *BinaryCIF* files, this is a dictionary
+            that can be encoded into *MessagePack*.
+        """
+        raise NotImplementedError()
+
+    def __str__(self):
+        return str(self.serialize())
+
+
+class _HierarchicalContainer(_Component, MutableMapping, metaclass=ABCMeta):
+    """
+    A container for hierarchical data in BinaryCIF files.
+    For example, the file contains multiple blocks, each block contains
+    multiple categories and each category contains multiple columns.
+
+    It uses lazy deserialization:
+    A component is only deserialized from the serialized data, if it
+    is accessed.
+    The deserialized component is then cached in the container.
+    """
+
+    def __init__(self, elements=None):
+        if elements is None:
+            elements = {}
+        for element in elements.values():
+            if not isinstance(element, (dict, self.subcomponent_class())):
+                raise TypeError(
+                    f"Expected '{self.subcomponent_class().__name__}', "
+                    f"but got '{type(element).__name__}'"
+                )
+        self._elements = elements
+
+    @staticmethod
+    def _deserialize_elements(content, take_key_from):
+        """
+        Lazily deserialize the elements of this container.
+
+        Parameters
+        ----------
+        content : dict
+            The serialized content describing the elements for this
+            container.
+        take_key_from : str
+            The key in each element of `content`, whose value is used as
+            the key for the respective element.
+
+        Returns
+        -------
+        elements : dict
+            The elements that should be stored in this container.
+            This return value can be given to the constructor.
+        """
+        elements = {}
+        for serialized_element in content:
+            key = serialized_element[take_key_from]
+            # Lazy deserialization
+            # -> keep serialized for now and deserialize later if needed
+            elements[key] = serialized_element
+        return elements
+
+    def _serialize_elements(self, store_key_in=None):
+        """
+        Serialize the elements that are stored in this container.
+
+        Each element that is still serialized (due to lazy
+        deserialization), is kept as it is.
+
+        Parameters
+        ----------
+        store_key_in: str, optional
+            If given, the key of each element is stored as value in the
+            serialized element.
+            This is basically the reverse operation of `take_key_from` in
+            :meth:`_deserialize_elements()`.
+        """
+        serialized_elements = []
+        for key, element in self._elements.items():
+            if isinstance(element, self.subcomponent_class()):
+                try:
+                    serialized_element = element.serialize()
+                except Exception:
+                    raise SerializationError(f"Failed to serialize element '{key}'")
+            else:
+                # Element is already stored in serialized form
+                serialized_element = element
+            if store_key_in is not None:
+                serialized_element[store_key_in] = key
+                serialized_elements.append(serialized_element)
+        return serialized_elements
+
+    def __getitem__(self, key):
+        element = self._elements[key]
+        if not isinstance(element, self.subcomponent_class()):
+            # Element is stored in serialized form
+            # -> must be deserialized first
+            try:
+                element = self.subcomponent_class().deserialize(element)
+            except Exception:
+                raise DeserializationError(f"Failed to deserialize element '{key}'")
+            # Update container with deserialized object
+            self._elements[key] = element
+        return element
+
+    def __setitem__(self, key, element):
+        if isinstance(element, self.subcomponent_class()):
+            pass
+        elif isinstance(element, _HierarchicalContainer):
+            # A common mistake may be to use the wrong container type
+            raise TypeError(
+                f"Expected '{self.subcomponent_class().__name__}', "
+                f"but got '{type(element).__name__}'"
+            )
+        else:
+            try:
+                element = self.subcomponent_class().deserialize(element)
+            except Exception:
+                raise DeserializationError("Failed to deserialize given value")
+        self._elements[key] = element
+
+    def __delitem__(self, key):
+        del self._elements[key]
+
+    # Implement `__contains__()` explicitly,
+    # because the mixin method unnecessarily deserializes the value, if available
+    def __contains__(self, key):
+        return key in self._elements
+
+    def __iter__(self):
+        return iter(self._elements)
+
+    def __len__(self):
+        return len(self._elements)
+
+    def __eq__(self, other):
+        if not isinstance(other, type(self)):
+            return False
+        if set(self.keys()) != set(other.keys()):
+            return False
+        for key in self.keys():
+            if self[key] != other[key]:
+                return False
+        return True

biotite/structure/io/pdbx/compress.py

@@ -0,0 +1,321 @@
+__all__ = ["compress"]
+__name__ = "biotite.structure.io.pdbx"
+__author__ = "Patrick Kunzmann"
+
+import itertools
+import msgpack
+import numpy as np
+import biotite.structure.io.pdbx.bcif as bcif
+from biotite.structure.io.pdbx.bcif import _encode_numpy as encode_numpy
+from biotite.structure.io.pdbx.encoding import (
+    ByteArrayEncoding,
+    DeltaEncoding,
+    FixedPointEncoding,
+    IntegerPackingEncoding,
+    RunLengthEncoding,
+    StringArrayEncoding,
+)
+
+
+def compress(data, float_tolerance=1e-6):
+    """
+    Try to reduce the size of a *BinaryCIF* file (or block, category, etc.) by testing
+    different data encodings for each data array and selecting the one, which results in
+    the smallest size.
+
+    Parameters
+    ----------
+    data : BinaryCIFFile or BinaryCIFBlock or BinaryCIFCategory or BinaryCIFColumn or BinaryCIFData
+        The data to compress.
+
+    Returns
+    -------
+    compressed_file : BinaryCIFFile or BinaryCIFBlock or BinaryCIFCategory or BinaryCIFColumn or BinaryCIFData
+        The compressed data with the same type as the input data.
+        If no improved compression is found for a :class:`BinaryCIFData` array,
+        the input data is kept.
+        Hence, the return value is no deep copy of the input data.
+    float_tolerance : float, optional
+        The relative error that is accepted when compressing floating point numbers.
+
+    Examples
+    --------
+
+    >>> from io import BytesIO
+    >>> pdbx_file = BinaryCIFFile()
+    >>> set_structure(pdbx_file, atom_array_stack)
+    >>> # Write uncompressed file
+    >>> uncompressed_file = BytesIO()
+    >>> pdbx_file.write(uncompressed_file)
+    >>> _ = uncompressed_file.seek(0)
+    >>> print(f"{len(uncompressed_file.read()) // 1000} KB")
+    927 KB
+    >>> # Write compressed file
+    >>> pdbx_file = compress(pdbx_file)
+    >>> compressed_file = BytesIO()
+    >>> pdbx_file.write(compressed_file)
+    >>> _ = compressed_file.seek(0)
+    >>> print(f"{len(compressed_file.read()) // 1000} KB")
+    111 KB
+    """
+    match type(data):
+        case bcif.BinaryCIFFile:
+            return _compress_file(data, float_tolerance)
+        case bcif.BinaryCIFBlock:
+            return _compress_block(data, float_tolerance)
+        case bcif.BinaryCIFCategory:
+            return _compress_category(data, float_tolerance)
+        case bcif.BinaryCIFColumn:
+            return _compress_column(data, float_tolerance)
+        case bcif.BinaryCIFData:
+            return _compress_data(data, float_tolerance)
+        case _:
+            raise TypeError(f"Unsupported type {type(data).__name__}")
+
+
+def _compress_file(bcif_file, float_tolerance):
+    compressed_file = bcif.BinaryCIFFile()
+    for block_name, bcif_block in bcif_file.items():
+        compressed_block = _compress_block(bcif_block, float_tolerance)
+        compressed_file[block_name] = compressed_block
+    return compressed_file
+
+
+def _compress_block(bcif_block, float_tolerance):
+    compressed_block = bcif.BinaryCIFBlock()
+    for category_name, bcif_category in bcif_block.items():
+        compressed_category = _compress_category(bcif_category, float_tolerance)
+        compressed_block[category_name] = compressed_category
+    return compressed_block
+
+
+def _compress_category(bcif_category, float_tolerance):
+    compressed_category = bcif.BinaryCIFCategory()
+    for column_name, bcif_column in bcif_category.items():
+        compressed_column = _compress_column(bcif_column, float_tolerance)
+        compressed_category[column_name] = compressed_column
+    return compressed_category
+
+
+def _compress_column(bcif_column, float_tolerance):
+    data = _compress_data(bcif_column.data, float_tolerance)
+    if bcif_column.mask is not None:
+        mask = _compress_data(bcif_column.mask, float_tolerance)
+    else:
+        mask = None
+    return bcif.BinaryCIFColumn(data, mask)
+
+
+def _compress_data(bcif_data, float_tolerance):
+    array = bcif_data.array
+    if len(array) == 1:
+        # No need to compress a single value -> Use default uncompressed encoding
+        return bcif.BinaryCIFData(array)
+
+    if np.issubdtype(array.dtype, np.str_):
+        # Leave encoding empty for now, as it is explicitly set later
+        encoding = StringArrayEncoding(data_encoding=[], offset_encoding=[])
+        # Run encode to initialize the data and offset arrays
+        indices = encoding.encode(array)
+        offsets = np.cumsum([0] + [len(s) for s in encoding.strings])
+        encoding.data_encoding, _ = _find_best_integer_compression(indices)
+        encoding.offset_encoding, _ = _find_best_integer_compression(offsets)
+        return bcif.BinaryCIFData(array, [encoding])
+
+    elif np.issubdtype(array.dtype, np.floating):
+        to_integer_encoding = FixedPointEncoding(
+            10 ** _get_decimal_places(array, float_tolerance)
+        )
+        integer_array = to_integer_encoding.encode(array)
+        best_encoding, size_compressed = _find_best_integer_compression(integer_array)
+        if size_compressed < _data_size_in_file(bcif.BinaryCIFData(array)):
+            return bcif.BinaryCIFData(array, [to_integer_encoding] + best_encoding)
+        else:
+            # The float array is smaller -> encode it directly as bytes
+            return bcif.BinaryCIFData(array, [ByteArrayEncoding()])
+
+    elif np.issubdtype(array.dtype, np.integer):
+        array = _to_smallest_integer_type(array)
+        encodings, _ = _find_best_integer_compression(array)
+        return bcif.BinaryCIFData(array, encodings)
+
+    else:
+        raise TypeError(f"Unsupported data type {array.dtype}")
+
+
+def _find_best_integer_compression(array):
+    """
+    Try different data encodings on an integer array and return the one that results in
+    the smallest size.
+    """
+    best_encoding_sequence = None
+    smallest_size = np.inf
+
+    for use_delta in [False, True]:
+        if use_delta:
+            encoding = DeltaEncoding()
+            array_after_delta = encoding.encode(array)
+            encodings_after_delta = [encoding]
+        else:
+            encodings_after_delta = []
+            array_after_delta = array
+        for use_run_length in [False, True]:
+            # Use encoded data from previous step to save time
+            if use_run_length:
+                encoding = RunLengthEncoding()
+                array_after_rle = encoding.encode(array_after_delta)
+                encodings_after_rle = encodings_after_delta + [encoding]
+            else:
+                encodings_after_rle = encodings_after_delta
+                array_after_rle = array_after_delta
+            for packed_byte_count in [None, 1, 2]:
+                if packed_byte_count is not None:
+                    # Quickly check this heuristic
+                    # to avoid computing an exploding packed data array
+                    if (
+                        _estimate_packed_length(array_after_rle, packed_byte_count)
+                        >= array_after_rle.nbytes
+                    ):
+                        # Packing would not reduce the size
+                        continue
+                    encoding = IntegerPackingEncoding(packed_byte_count)
+                    array_after_packing = encoding.encode(array_after_rle)
+                    encodings_after_packing = encodings_after_rle + [encoding]
+                else:
+                    encodings_after_packing = encodings_after_rle
+                    array_after_packing = array_after_rle
+                encoding = ByteArrayEncoding()
+                encoded_array = encoding.encode(array_after_packing)
+                encodings = encodings_after_packing + [encoding]
+                # Pack data directly instead of using the BinaryCIFData class
+                # to avoid the unnecessary re-encoding of the array,
+                # as it is already available in 'encoded_array'
+                serialized_encoding = [enc.serialize() for enc in encodings]
+                serialized_data = {
+                    "data": encoded_array,
+                    "encoding": serialized_encoding,
+                }
+                size = _data_size_in_file(serialized_data)
+                if size < smallest_size:
+                    best_encoding_sequence = encodings
+                    smallest_size = size
+    return best_encoding_sequence, smallest_size
+
+
+def _estimate_packed_length(array, packed_byte_count):
+    """
+    Estimate the length of an integer array after packing it with a given number of
+    bytes.
+
+    Parameters
+    ----------
+    array : numpy.ndarray
+        The array to pack.
+    packed_byte_count : int
+        The number of bytes used for packing.
+
+    Returns
+    -------
+    length : int
+        The estimated length of the packed array.
+    """
+    # Use int64 to avoid integer overflow in the following line
+    max_val_per_element = np.int64(2 ** (8 * packed_byte_count))
+    n_bytes_per_element = packed_byte_count * (np.abs(array // max_val_per_element) + 1)
+    return np.sum(n_bytes_per_element, dtype=np.int64)
+
+
+def _to_smallest_integer_type(array):
+    """
+    Convert an integer array to the smallest possible integer type, that is still able
+    to represent all values in the array.
+
+    Parameters
+    ----------
+    array : numpy.ndarray
+        The array to convert.
+
+    Returns
+    -------
+    array : numpy.ndarray
+        The converted array.
+    """
+    if array.min() >= 0:
+        for dtype in [np.uint8, np.uint16, np.uint32, np.uint64]:
+            if np.all(array <= np.iinfo(dtype).max):
+                return array.astype(dtype)
+    for dtype in [np.int8, np.int16, np.int32, np.int64]:
+        if np.all(array >= np.iinfo(dtype).min) and np.all(
+            array <= np.iinfo(dtype).max
+        ):
+            return array.astype(dtype)
+    raise ValueError("Array is out of bounds for all integer types")
+
+
+def _data_size_in_file(data):
+    """
+    Get the size of the data, it would have when written into a *BinaryCIF* file.
+
+    Parameters
+    ----------
+    data : BinaryCIFData or dict
+        The data array whose size is measured.
+        Can be either a :class:`BinaryCIFData` object or already serialized data.
+
+    Returns
+    -------
+    size : int
+        The size of the data array in the file in bytes.
+    """
+    if isinstance(data, bcif.BinaryCIFData):
+        data = data.serialize()
+    bytes_in_file = msgpack.packb(data, use_bin_type=True, default=encode_numpy)
+    return len(bytes_in_file)
+
+
+def _get_decimal_places(array, tol):
+    """
+    Get the number of decimal places in a floating point array.
+
+    Parameters
+    ----------
+    array : numpy.ndarray
+        The array to analyze.
+    tol : float, optional
+        The relative tolerance allowed when the values are cut off after the returned
+        number of decimal places.
+
+    Returns
+    -------
+    decimals : int
+        The number of decimal places.
+    """
+    # Decimals of NaN or infinite values do not make sense
+    # and 0 would give NaN when rounding on decimals
+    array = array[np.isfinite(array) & (array != 0)]
+    for decimals in itertools.count(start=-_order_magnitude(array)):
+        error = np.abs(np.round(array, decimals) - array)
+        if np.all(error < tol * np.abs(array)):
+            return decimals
+
+
+def _order_magnitude(array):
+    """
+    Get the order of magnitude of floating point values.
+
+    Parameters
+    ----------
+    array : ndarray, dtype=float
+        The value to analyze.
+
+    Returns
+    -------
+    magnitude : int
+        The order of magnitude, i.e. the maximum exponent a number in the array would
+        have in scientific notation, if only one digit is left of the decimal point.
+    """
+    array = array[array != 0]
+    if len(array) == 0:
+        # No non-zero values -> define order of magnitude as 0
+        return 0
+    return int(np.max(np.floor(np.log10(np.abs(array)))).item())