scoring.py

"""Objects for defining and customizing the target cost of a contraction."""

import re
import math
import functools

# the default weighting for comparing flops vs mops
DEFAULT_COMBO_FACTOR = 64


class Objective:
    """Base mixin class for all objectives."""

    __slots__ = ()

    def __call__(self, trial):
        """The core method that takes a ``trial`` generated by a contraction
        path driver and scores it to report to a hyper-optimizer. It might also
        update the parameters in the trial to reflect the desired cost.
        """
        raise NotImplementedError

    def __repr__(self):
        params = {k: getattr(self, k) for k in getattr(self, "__slots__", ())}
        return (
            f"{self.__class__.__name__}("
            + ", ".join(f"{k}={v}" for k, v in params.items())
            + ")"
        )

    def __hash__(self):
        return hash(repr(self))


# ------------------------ exact contraction scoring ------------------------ #


def ensure_basic_quantities_are_computed(trial):
    if not all(q in trial for q in ("flops", "write", "size")):
        stats = trial["tree"].contract_stats()

        if "flops" not in trial:
            trial["flops"] = stats["flops"]
        if "write" not in trial:
            trial["write"] = stats["write"]
        if "size" not in trial:
            trial["size"] = stats["size"]


class ExactObjective(Objective):
    """Mixin class for all exact objectives."""

    def cost_local_tree_node(self, tree, node):
        """The cost of a single ``node`` in ``tree``, according to this
        objective. Used for subtree reconfiguration.
        """
        raise NotImplementedError

    def score_local(self, **kwargs):
        """The score to give a single contraction, according to the given
        ``kwargs``. Used in ``simulated_anneal``.
        """
        raise NotImplementedError

    def score_slice_index(self, costs, ix):
        """The score to give possibly slicing ``ix``, according to
        the given ``costs``. Used in the ``SliceFinder`` optimization.
        """
        raise NotImplementedError

    def get_dynamic_programming_minimize(self):
        """Get the argument for optimal optimization, used in
        subtree reconfiguration.
        """
        raise NotImplementedError


class FlopsObjective(ExactObjective):
    """Objective that scores based on estimated floating point operations.

    Parameters
    ----------
    secondary_weight : float, optional
        Weighting factor for secondary objectives (max size and total write).
        Default is 1e-3.
    """

    __slots__ = ("secondary_weight",)

    def __init__(self, secondary_weight=1e-3):
        self.secondary_weight = secondary_weight
        super().__init__()

    def cost_local_tree_node(self, tree, node):
        return tree.get_flops(node)

    def score_local(self, **kwargs):
        f = kwargs["flops"]
        try:
            # accept iterables
            f = sum(f)
        except TypeError:
            pass
        return math.log2(f)

    def score_slice_index(self, costs, ix):
        return math.log(
            costs._flop_reductions[ix]
            + costs._write_reductions[ix] * self.secondary_weight
            + 1
        )

    def get_dynamic_programming_minimize(self):
        return "flops"

    def __call__(self, trial):
        ensure_basic_quantities_are_computed(trial)
        return (
            math.log2(trial["flops"])
            + self.secondary_weight * math.log2(trial["write"])
            + self.secondary_weight * math.log2(trial["size"])
        )


class WriteObjective(ExactObjective):
    """Objective that scores based on estimated total write, i.e. the sum of
    sizes of all intermediates. This is relevant for completely memory-bound
    contractions, and also for back-propagation.

    Parameters
    ----------
    secondary_weight : float, optional
        Weighting factor for secondary objectives (max size and total flops).
        Default is 1e-3.
    """

    __slots__ = ("secondary_weight",)

    def __init__(self, secondary_weight=1e-3):
        self.secondary_weight = secondary_weight
        super().__init__()

    def cost_local_tree_node(self, tree, node):
        return tree.get_size(node)

    def score_local(self, **kwargs):
        s = kwargs["size"]
        try:
            # accept iterables
            s = sum(s)
        except TypeError:
            pass
        return math.log2(s)

    def score_slice_index(self, costs, ix):
        return math.log(
            costs._flop_reductions[ix] * self.secondary_weight
            + costs._write_reductions[ix]
            + 1
        )

    def get_dynamic_programming_minimize(self):
        return "write"

    def __call__(self, trial):
        ensure_basic_quantities_are_computed(trial)
        return (
            +self.secondary_weight * math.log2(trial["flops"])
            + math.log2(trial["write"])
            + self.secondary_weight * math.log2(trial["size"])
        )


class SizeObjective(ExactObjective):
    """Objective that scores based on maximum intermediate size.

    Parameters
    ----------
    secondary_weight : float, optional
        Weighting factor for secondary objectives (total flops and total
        write). Default is 1e-3.
    """

    __slots__ = ("secondary_weight",)

    def __init__(self, secondary_weight=1e-3):
        self.secondary_weight = secondary_weight
        super().__init__()

    def cost_local_tree_node(self, tree, node):
        return tree.get_size(node)

    def score_local(self, **kwargs):
        s = kwargs["size"]
        try:
            # accept iterables
            s = max(s)
        except TypeError:
            pass
        return math.log2(s)

    def score_slice_index(self, costs, ix):
        return math.log(
            costs._flop_reductions[ix] * self.secondary_weight
            + costs._write_reductions[ix]
            + 1
        )

    def get_dynamic_programming_minimize(self):
        return "size"

    def __call__(self, trial):
        ensure_basic_quantities_are_computed(trial)
        return (
            +self.secondary_weight * math.log2(trial["flops"])
            + self.secondary_weight * math.log2(trial["write"])
            + math.log2(trial["size"])
        )


class ComboObjective(ExactObjective):
    """Objective that scores based on a combination of estimated floating point
    operations and total write, according to:

    .. math::

        \\log_2(\\text{flops} + \\alpha \\times \\text{write})

    Where alpha is the ``factor`` parameter of this objective, that describes
    approximately how much slower write speeds are.

    Parameters
    ----------
    factor : float, optional
        Weighting factor for total write. Default is 64.
    """

    __slots__ = ("factor",)

    def __init__(
        self,
        factor=DEFAULT_COMBO_FACTOR,
    ):
        self.factor = factor
        super().__init__()

    def cost_local_tree_node(self, tree, node):
        return tree.get_flops(node) + self.factor * tree.get_size(node)

    def score_local(self, **kwargs):
        f = kwargs["flops"]
        try:
            f = sum(f)
        except TypeError:
            f = ()
        try:
            # accept iterables
            f = sum(f)
        except TypeError:
            pass
        w = kwargs["size"]
        try:
            # accept iterables
            w = sum(w)
        except TypeError:
            pass
        return math.log2(f + self.factor * w)

    def score_slice_index(self, costs, ix):
        return math.log(
            costs._flop_reductions[ix]
            + costs._write_reductions[ix] * self.factor
            + 1
        )

    def get_dynamic_programming_minimize(self):
        return f"combo-{self.factor}"

    def __call__(self, trial):
        ensure_basic_quantities_are_computed(trial)
        return math.log2(trial["flops"] + self.factor * trial["write"])


class LimitObjective(ExactObjective):
    """Objective that scores based on a maximum of either estimated floating
    point operations or the total write, weighted by some factor:

    .. math::


        \\sum_{c}
        max(\\text{flops}_i, \\alpha \\times \\text{write}_i)

    For each contraction $i$. Where alpha is the ``factor`` parameter of this
    objective, that describes approximately how much slower write speeds are.
    This assumes that one or the other is the limiting factor.

    Parameters
    ----------
    factor : float, optional
        Weighting factor for total write. Default is 64.
    """

    def __init__(self, factor=DEFAULT_COMBO_FACTOR):
        self.factor = factor
        super().__init__()

    def cost_local_tree_node(self, tree, node):
        return max(tree.get_flops(node), self.factor * tree.get_size(node))

    def score_local(self, **kwargs):
        f = kwargs["flops"]
        w = kwargs["size"]
        try:
            return math.log2(
                sum(max(fi, self.factor * wi) for fi, wi in zip(f, w))
            )
        except TypeError:
            return math.log2(max(f, self.factor * w))

    def score_slice_index(self, costs, ix):
        return math.log(
            costs._flop_reductions[ix]
            + costs._write_reductions[ix] * self.factor
            + 1
        )

    def get_dynamic_programming_minimize(self):
        return f"limit-{self.factor}"

    def __call__(self, trial):
        tree = trial["tree"]
        return math.log2(tree.combo_cost(factor=self.factor, combine=max))


# --------------------- compressed contraction scoring ---------------------- #


class CompressedStatsTracker:
    __slots__ = (
        "chi",
        "flops",
        "max_size",
        "peak_size",
        "write",
        "total_size",
        "total_size_post_contract",
        "contracted_size",
        "size_change",
        "flops_change",
    )

    def __init__(self, hg, chi):
        if chi == "auto":
            self.chi = max(hg.size_dict.values()) ** 2
        else:
            self.chi = chi

        # local params -> don't depend on history
        self.total_size = 0
        self.total_size_post_contract = 0
        self.contracted_size = 0
        self.size_change = 0
        self.flops_change = 0

        # global params -> depend on history
        self.flops = 0
        self.max_size = 0

        # initial tensors contribute to size
        for i in hg.nodes:
            sz_i = hg.node_size(i)
            self.max_size = max(self.max_size, sz_i)
            self.total_size += sz_i

        self.write = self.peak_size = self.total_size

    def copy(self):
        new = object.__new__(self.__class__)
        for attr in self.__slots__:
            setattr(new, attr, getattr(self, attr))
        return new

    def update_pre_step(self):
        self.size_change = 0
        self.flops_change = 0

    def update_pre_compress(self, hg, *nodes):
        # subtract tensors size and also their neighbors size (since both will
        # change with compression)
        self.size_change -= hg.neighborhood_size(nodes)
        self.flops_change += hg.neighborhood_compress_cost(self.chi, nodes)

    def update_post_compress(self, hg, *nodes):
        # add new tensors size and also its neighbors (since these will have
        # changed with compression)
        self.size_change += hg.neighborhood_size(nodes)

    def update_pre_contract(self, hg, i, j):
        # remove pair of tensors from size
        self.size_change -= hg.node_size(i) + hg.node_size(j)
        # add flops of just the contraction
        self.flops_change += hg.contract_pair_cost(i, j)

    def update_post_contract(self, hg, ij):
        self.contracted_size = hg.node_size(ij)
        self.size_change += self.contracted_size

        # compute here before potential compressions:
        # the peak total size of concurrent intermediates
        self.total_size_post_contract = self.total_size + self.size_change

    def update_post_step(self):
        self.max_size = max(self.max_size, self.contracted_size)
        self.peak_size = max(self.peak_size, self.total_size_post_contract)
        self.total_size += self.size_change
        self.flops += self.flops_change
        self.write += self.contracted_size

    def update_score(self, other):
        self.flops = other.flops + self.flops_change
        self.write = other.write + self.contracted_size
        self.max_size = max(other.max_size, self.contracted_size)
        self.peak_size = max(other.peak_size, self.total_size_post_contract)

        if self.max_size > self.peak_size:
            raise RuntimeError(
                f"max_size={self.max_size} > peak_size={self.peak_size}"
            )

    @property
    def combo_score(self):
        return math.log2(self.flops + DEFAULT_COMBO_FACTOR * self.write + 1)

    @property
    def score(self):
        raise NotImplementedError

    def describe(self, join=" "):
        F = math.log10(max(1, self.flops))
        C = math.log10(
            max(
                1,
                self.flops
                + getattr(self, "factor", DEFAULT_COMBO_FACTOR) * self.write,
            )
        )
        S = math.log2(max(1, self.max_size))
        P = math.log2(max(1, self.peak_size))
        return join.join((
            f"F={F:.2f}", f"C={C:.2f}", f"S={S:.2f}", f"P={P:.2f}"
        ))

    def __repr__(self):\
        return (
            f"<{self.__class__.__name__}({self.describe(join=', ')})>"
        )


class CompressedStatsTrackerSize(CompressedStatsTracker):
    __slots__ = CompressedStatsTracker.__slots__ + ("secondary_weight",)

    def __init__(self, hg, chi, secondary_weight=1e-3):
        self.secondary_weight = secondary_weight
        super().__init__(hg, chi)

    @property
    def score(self):
        return (
            math.log2(self.max_size)
            + math.log2(self.flops + 1) * self.secondary_weight
        )


class CompressedStatsTrackerPeak(CompressedStatsTracker):
    __slots__ = CompressedStatsTracker.__slots__ + ("secondary_weight",)

    def __init__(self, hg, chi, secondary_weight=1e-3):
        self.secondary_weight = secondary_weight
        super().__init__(hg, chi)

    @property
    def score(self):
        return (
            math.log2(self.peak_size)
            + math.log2(self.flops + 1) * self.secondary_weight
        )


class CompressedStatsTrackerWrite(CompressedStatsTracker):
    __slots__ = CompressedStatsTracker.__slots__ + ("secondary_weight",)

    def __init__(self, hg, chi, secondary_weight=1e-3):
        self.secondary_weight = secondary_weight
        super().__init__(hg, chi)

    @property
    def score(self):
        return (
            math.log2(self.write)
            + math.log2(self.flops + 1) * self.secondary_weight
        )


class CompressedStatsTrackerFlops(CompressedStatsTracker):
    __slots__ = CompressedStatsTracker.__slots__ + ("secondary_weight",)

    def __init__(self, hg, chi, secondary_weight=1e-3):
        self.secondary_weight = secondary_weight
        super().__init__(hg, chi)

    @property
    def score(self):
        return (
            math.log10(self.flops + 1)
            + math.log10(self.peak_size) * self.secondary_weight
        )


class CompressedStatsTrackerCombo(CompressedStatsTracker):
    __slots__ = CompressedStatsTracker.__slots__ + ("factor",)

    def __init__(self, hg, chi, factor=DEFAULT_COMBO_FACTOR):
        self.factor = factor
        super().__init__(hg, chi)

    @property
    def score(self):
        return math.log2(self.flops + self.factor * self.write + 1)


class CompressedObjective(Objective):
    """Mixin for objectives that score based on a compressed contraction."""

    def __init__(self, chi="auto", compress_late=False):
        self.chi = chi
        self.compress_late = compress_late
        super().__init__()

    def get_compressed_stats_tracker(self, hg):
        """Return a tracker for compressed contraction stats.

        Parameters
        ----------
        hg : Hypergraph
            The hypergraph to track stats for.

        Returns
        -------
        CompressedStatsTracker
            The tracker.
        """
        raise NotImplementedError

    def compute_compressed_stats(self, trial):
        tree = trial["tree"]
        if self.chi == "auto":
            chi = max(tree.size_dict.values()) ** 2
        else:
            chi = self.chi

        return tree.compressed_contract_stats(
            chi,
            compress_late=self.compress_late,
        )


class CompressedSizeObjective(CompressedObjective):
    """Objective that scores based on the maximum size intermediate tensor
    during a compressed contraction with maximum bond dimension ``chi``.

    Parameters
    ----------
    chi : int, optional
        Maximum bond dimension to use for the compressed contraction. Default
        is ``"auto"``, which will use the square of the maximum size of any
        input tensor dimension.
    compress_late : bool, optional
        Whether to compress the neighboring tensors just after (early) or just
        before (late) contracting tensors. Default is False, i.e. early.
    secondary_weight : float, optional
        Weighting factor for secondary objectives (flops and write).
        Default is 1e-3.
    """

    __slots__ = ("chi", "compress_late", "secondary_weight")

    def __init__(
        self,
        chi="auto",
        compress_late=False,
        secondary_weight=1e-3,
    ):
        self.secondary_weight = secondary_weight
        super().__init__(chi=chi, compress_late=compress_late)

    def get_compressed_stats_tracker(self, hg):
        return CompressedStatsTrackerSize(
            hg, self.chi, secondary_weight=self.secondary_weight
        )

    def __call__(self, trial):
        stats = self.compute_compressed_stats(trial)
        cr = (
            math.log2(stats.max_size)
            + self.secondary_weight * math.log2(stats.flops)
            + self.secondary_weight * math.log2(stats.write)
        )

        # overwrite stats with compressed versions
        trial["size"] = stats.max_size
        trial["flops"] = stats.flops
        trial["write"] = stats.write
        return cr


class CompressedPeakObjective(CompressedObjective):
    """Objective that scores based on the peak total concurrent size of
    intermediate tensors during a compressed contraction with maximum bond
    dimension ``chi``.

    Parameters
    ----------
    chi : int, optional
        Maximum bond dimension to use for the compressed contraction. Default
        is ``"auto"``, which will use the square of the maximum size of any
        input tensor dimension.
    compress_late : bool, optional
        Whether to compress the neighboring tensors just after (early) or just
        before (late) contracting tensors. Default is False, i.e. early.
    secondary_weight : float, optional
        Weighting factor for secondary objectives (flops and write).
        Default is 1e-3.
    """

    __slots__ = ("chi", "compress_late", "secondary_weight")

    def __init__(
        self,
        chi="auto",
        compress_late=False,
        secondary_weight=1e-3,
    ):
        self.secondary_weight = secondary_weight
        super().__init__(chi=chi, compress_late=compress_late)

    def get_compressed_stats_tracker(self, hg):
        return CompressedStatsTrackerPeak(
            hg, chi=self.chi, secondary_weight=self.secondary_weight
        )

    def __call__(self, trial):
        stats = self.compute_compressed_stats(trial)
        cr = (
            math.log2(stats.peak_size)
            + self.secondary_weight * math.log2(stats.flops)
            + self.secondary_weight * math.log2(stats.write)
        )

        # overwrite stats with compressed versions
        trial["size"] = stats.peak_size
        trial["flops"] = stats.flops
        trial["write"] = stats.write
        return cr


class CompressedWriteObjective(CompressedObjective):
    """Objective that scores based on the total cumulative size of
    intermediate tensors during a compressed contraction with maximum bond
    dimension ``chi``.

    Parameters
    ----------
    chi : int, optional
        Maximum bond dimension to use for the compressed contraction. Default
        is ``"auto"``, which will use the square of the maximum size of any
        input tensor dimension.
    compress_late : bool, optional
        Whether to compress the neighboring tensors just after (early) or just
        before (late) contracting tensors. Default is False, i.e. early.
    secondary_weight : float, optional
        Weighting factor for secondary objectives (flops and peak size).
        Default is 1e-3.
    """

    __slots__ = ("chi", "compress_late", "secondary_weight")

    def __init__(
        self,
        chi="auto",
        compress_late=False,
        secondary_weight=1e-3,
    ):
        self.secondary_weight = secondary_weight
        super().__init__(chi=chi, compress_late=compress_late)

    def get_compressed_stats_tracker(self, hg):
        return CompressedStatsTrackerWrite(
            hg, chi=self.chi, secondary_weight=self.secondary_weight
        )

    def __call__(self, trial):
        stats = self.compute_compressed_stats(trial)
        cr = (
            math.log2(stats.write)
            + self.secondary_weight * math.log2(stats.flops)
            + self.secondary_weight * math.log2(stats.peak_size)
        )

        # overwrite stats with compressed versions
        trial["size"] = stats.write
        trial["flops"] = stats.flops
        trial["write"] = stats.write
        return cr


class CompressedFlopsObjective(CompressedObjective):
    """Objective that scores based on the total contraction flops
    intermediate tensors during a compressed contraction with maximum bond
    dimension ``chi``.

    Parameters
    ----------
    chi : int, optional
        Maximum bond dimension to use for the compressed contraction. Default
        is ``"auto"``, which will use the square of the maximum size of any
        input tensor dimension.
    compress_late : bool, optional
        Whether to compress the neighboring tensors just after (early) or just
        before (late) contracting tensors. Default is False, i.e. early.
    secondary_weight : float, optional
        Weighting factor for secondary objectives (write and peak size).
        Default is 1e-3.
    """

    __slots__ = ("chi", "compress_late", "secondary_weight")

    def __init__(
        self,
        chi="auto",
        compress_late=False,
        secondary_weight=1e-3,
    ):
        self.secondary_weight = secondary_weight
        super().__init__(chi=chi, compress_late=compress_late)

    def get_compressed_stats_tracker(self, hg):
        return CompressedStatsTrackerFlops(
            hg, chi=self.chi, secondary_weight=self.secondary_weight
        )

    def __call__(self, trial):
        stats = self.compute_compressed_stats(trial)
        cr = (
            math.log2(stats.flops)
            + self.secondary_weight * math.log2(stats.write)
            + self.secondary_weight * math.log2(stats.peak_size)
        )

        # overwrite stats with compressed versions
        trial["size"] = stats.max_size
        trial["flops"] = stats.flops
        trial["write"] = stats.write
        return cr


class CompressedComboObjective(CompressedObjective):
    __slots__ = ("chi", "compress_late", "factor")

    def __init__(
        self,
        chi="auto",
        compress_late=False,
        factor=DEFAULT_COMBO_FACTOR,
    ):
        self.factor = factor
        super().__init__(chi=chi, compress_late=compress_late)

    def get_compressed_stats_tracker(self, hg):
        return CompressedStatsTrackerCombo(
            hg, chi=self.chi, factor=self.factor
        )

    def __call__(self, trial):
        stats = self.compute_compressed_stats(trial)

        flops = stats.flops
        write = stats.write

        cr = math.log2(flops + self.factor * write)

        # overwrite stats with compressed versions
        trial["size"] = stats.max_size
        trial["flops"] = flops
        trial["write"] = write
        return cr


score_matcher = re.compile(
    # exact scoring functions
    r"("
    r"flops|"
    r"size|"
    r"write|"
    r"combo|"
    r"limit|"
    # compressed scoring functions
    r"flops-compressed|"
    r"size-compressed|"
    r"max-compressed|"
    r"peak-compressed|"
    r"write-compressed|"
    r"combo-compressed"
    r")-*(\d*)"
)


def parse_minimize(minimize):
    match = score_matcher.fullmatch(minimize)
    if not match:
        raise ValueError(f"No score function '{minimize}' found.")

    which, param = match.groups()
    return which, param


@functools.lru_cache(maxsize=128)
def _get_score_fn_str_cached(minimize):
    which, param = parse_minimize(minimize)

    if which == "flops":
        return FlopsObjective()

    if which == "write":
        return WriteObjective()

    if which == "size":
        return SizeObjective()

    if which == "combo":
        factor = float(param) if param else DEFAULT_COMBO_FACTOR
        return ComboObjective(factor=factor)

    if which == "limit":
        factor = float(param) if param else DEFAULT_COMBO_FACTOR
        return LimitObjective(factor=factor)

    if which in ("max-compressed", "size-compressed"):
        chi = int(param) if param else "auto"
        return CompressedSizeObjective(chi=chi)

    if which == "peak-compressed":
        chi = int(param) if param else "auto"
        return CompressedPeakObjective(chi=chi)

    if which == "write-compressed":
        chi = int(param) if param else "auto"
        return CompressedWriteObjective(chi=chi)

    if which == "flops-compressed":
        chi = int(param) if param else "auto"
        return CompressedFlopsObjective(chi=chi)

    if which == "combo-compressed":
        chi = int(param) if param else "auto"
        return CompressedComboObjective(chi=chi)

    raise ValueError(f"No objective function named '{minimize}' found.")


def get_score_fn(minimize):
    if isinstance(minimize, str):
        return _get_score_fn_str_cached(minimize)
    if callable(minimize):
        # custom objective function
        return minimize
    raise TypeError("minimize must be a string or callable.")
# ----------------------- multi-contraction scoring ------------------------- #


class MultiObjective(Objective):

    __slots__ = ("num_configs",)

    def __init__(self, num_configs):
        self.num_configs = num_configs

    def compute_mult(self, dims):
        raise NotImplementedError

    def estimate_node_mult(self, tree, node):
        return self.compute_mult([
            tree.size_dict[ix] for ix in tree.get_node_var_inds(node)
        ])

    def estimate_node_cache_mult(self, tree, node, sliced_ind_ordering):
        node_var_inds = tree.get_node_var_inds(node)

        # indices which are the first 'k' in the sliced ordering
        non_heavy_inds = [
            ix
            for ix in tree.get_node_var_inds(node)
            if ix not in sliced_ind_ordering[: len(node_var_inds)]
        ]

        # each of these cycles 'out of sync' and thus must be kept
        return self.compute_mult([tree.size_dict[ix] for ix in non_heavy_inds])


class MultiObjectiveDense(MultiObjective):
    """Number of intermediate configurations is expected to scale as if all
    configurations are present.
    """
    __slots__ = ("num_configs",)

    def compute_mult(self, dims):
        return math.prod(dims)


def expected_coupons(num_sub, num_total):
    """If we draw a random 'coupon` which can take `num_sub` different values
    `num_total` times, how many unique coupons will we expect?
    """
    return num_sub * (1 - (1 - 1 / num_sub) ** num_total)


class MultiObjectiveUniform(MultiObjective):
    """Number of intermediate configurations is expected to scale as if all
    configurations are randomly draw from a uniform distribution.
    """
    __slots__ = ("num_configs",)

    def compute_mult(self, dims):
        return expected_coupons(math.prod(dims), self.num_configs)


class MultiObjectiveLinear(MultiObjective):
    """Number of intermediate configurations is expected to scale linearly with
    respect to number of variable indices (e.g. VMC like 'locally connected'
    configurations).
    """
    __slots__ = ("num_configs", "coeff")

    def __init__(self, num_configs, coeff=1):
        self.coeff = coeff
        super().__init__(num_configs=num_configs)

    def compute_mult(self, dims):
        return min(self.coeff * len(dims), self.num_configs)