Source code for schrodinger.infra.mmbitset

"""
A pythonic wrapper for the mmbs library, providing access to bitsets.

The default error handler for the Bitset module is set to mm.error_handler.
If no error handler is specified for error_handler arguments, its value is
what will be used.

Copyright Schrodinger, LLC. All rights reserved.

"""
from schrodinger.infra import mm
from schrodinger.infra import structure


class Bitset:
    """
    This class represents a bitset, with each element either 1 or 0.

    A bitset is similar conceptually to a Python dictionary that contains
    keys that are integers from 1 to the length of the bitset, and the value
    for each key is an on/off boolean (1 or 0).

    Bitsets are usually used when keeping track of which atoms in the ct an
    operation needs to be performed on; in which case the length of the Bitset
    is the number of atoms in the ct.

    An instance of Bitset can be converted to a list containing the elements
    that are on (i.e. set to 1) via list(bitset) and to a set via set(bitset).

    This is an object-oriented wrapper for the underlying MMBS library. All
    state is stored in the C library. References to instances of this class
    can be used in direct mmbs library calls, as it will converted to the
    integer handle when necessary.

    """

    def __init__(self,
                 handle=None,
                 manage_handle=True,
                 error_handler=None,
                 size=None):
        """
        Initialize an object with an existing MMBS handle or a size.

        By default, the MMBS resources will be managed by the object. To
        keep these from being cleaned up on object deletion, set
        manage_handle=False.

        :param size: The size to use for a newly created bitset.
        :type size: int
        """

        if handle is None:
            if size is None:
                raise ValueError(
                    "Bitset constructor must provide an existing mmbs handle or a size specification."
                )
            self.handle = structure.MM_Bitset.createFromSize(size)
        elif isinstance(handle, structure.MM_Bitset):
            self.handle = handle
        else:
            self.handle = structure.MM_Bitset(handle)
        if not manage_handle:
            self.handle.release()

    @classmethod
    def from_list(cls, size, on_list):
        """
        Alternative constructor; returns a Bitset object initialized from a
        Python list.

        Example usage::

            bs = Bitset.from_list(st.atom_total, selected_atoms)

        where <st> is the structure object, and <selected_atoms> is a list of
        atom indices.

        :param size: The size to use for a newly created bitset.
        :type size: int

        :param on_list: Turn on the bits that are present in this iterable.
                Each item should be an int between 1 and <size>.
        """
        bs = cls(size=size)
        for index in on_list:
            bs.set(index)
        return bs

    def size(self):
        """
        Return the capacity (i.e. maximum index) of the Bitset.

        """
        return len(self.handle)

    def resize(self, size):
        """
        Modify the capacity of the Bitset.

        """
        self.handle.resize(size)

    def count(self):
        """
        Return the number of index values that are set to 'on'.

        """
        return self.handle.count()

    def get(self, index):
        """
        Get the value of the index.

        Returns True if the index value is on, False if it is off.

        """
        return self.handle[int(index)]

    def set(self, index):
        """
        Set the index value to 'on'.

        """
        self.handle[int(index)] = True

    def unset(self, index):
        """
        Set the index value to 'off'.

        """
        self.handle[int(index)] = False

    def invert(self):
        """Invert the bits in the set."""
        mm.mmbs_invert(self.handle)

    def range(self, start, finish):
        """
        Set all bits in a range from "start" to "finish"
        """
        mm.mmbs_range(self.handle, start, finish)

    def fill(self):
        """
        Set all bits to on.

        """
        self.handle.set()

    def all(self):
        """Are all bits on?"""
        return bool(self.handle.all())

    def copy(self):
        """Create another bitset that is a duplicate of this one."""
        new_handle = mm.mmbs_duplicate(self.handle)
        return self.__class__(new_handle)

    def __repr__(self):
        """
        Return a Python representation string for this bitset.

        """
        return "Bitset(%d)" % self.handle.getHandle()

    def __str__(self):
        """
        Return a string of this bitset.

        """
        return str(list(self))

    def __iter__(self):
        """
        Return an iterator for this bitset.
        Loops over all of the ON bits.

        """
        # don't return iter(self.handle) in case this is a temporary. If this
        # is cleaned up before the iterator, the iterator is invalid.
        yield from self.handle

    def __len__(self):
        """
        Returns the size of the bitset (Number of ON and OFF bits)
        """
        return self.size()

    def __index__(self):
        """Get the MMBs C handle"""
        return int(self.handle)

    def __getitem__(self, index):
        return self.handle[index]

    def __setitem__(self, index, value):
        self.handle[index] = value