Germ Selection Gate Penalty

[coreyostrove]

We recently ran into a scenario where it was desirable to perform germ selection in a manner which minimized the number of applications of a particularly expensive gate operation. To support that sort of use case this PR adds a new algorithm_kwarg option called gate_penalty for find_germs which is compatible with the 'greedy' and 'grasp' search algorithms which allows a user to specify and additional penalty factors for the number of instances of particular gates in a candidate germ set. Also included are some additional unit tests for both this new gate penalty as well as tests for the existing op_penalty option (which penalized the total number of gate operations overall).

[rileyjmurray]

Left several comments. Happy to iterate on them a bit!

[rileyjmurray]

Type annotation please :)

[rileyjmurray]

There are seven exact matches for this text in this file. This will make long-term maintainability difficult. I can see two or three ways around it.

Method 1: docstring injection decorators.

We can define a docstring-injection decorator

def set_docstring(docstr):
    def assign(fn):
        fn.__doc__ = docstr
        return fn
    return assign

and a module-wide constant

GATE_PENALTY_DESCRIPTION = \
  """
  gate_penalty : dict, optional (default None)
          An optional dictionary allowing the specification of gate-specific penalties to add for each instance
          of the specified gate(s) in each germ. Should be specified as a dictionary whose keys are strings
          corresponding to gate names, and whose values are the penalty factor to add for each instance of that
          gate. E.g. {'Gcnot':2} would correspond to a penalty term where each instance of a 'Gcnot' gate
          gets and additional 2 units added to the cost function for a candidate germ.
  """

then set docstrings as follows. The downside of this approach is that the docstring now has to come before the function signature, rather than after.

@set_docstring(
f"""
Calculate the score of a germ set with respect to a model.

More precisely, this function computes the maximum score (roughly equal
to the number of amplified parameters) for a cloud of models.
If `target_model` is given, it serves as the center of the cloud,
otherwise the cloud must be supplied directly via `neighborhood`.

Parameters
----------
germs : list
    The germ set

target_model : Model, optional
    The target model, used to generate a neighborhood of randomized models.

neighborhood : list of Models, optional
    The "cloud" of models for which scores are computed.  If not None, this
    overrides `target_model`, `neighborhood_size`, and `randomization_strength`.

neighborhood_size : int, optional
    Number of randomized models to construct around `target_model`.

randomization_strength : float, optional
    Strength of unitary randomizations, as passed to :meth:`target_model.randomize_with_unitary`.

score_func : {'all', 'worst'}
    Sets the objective function for scoring the eigenvalues. If 'all',
    score is ``sum(1/input_array)``. If 'worst', score is ``1/min(input_array)``.

op_penalty : float, optional
    Coefficient for a penalty linear in the sum of the germ lengths.

l1_penalty : float, optional
    Coefficient for a penalty linear in the number of germs.

num_nongauge_params : int, optional
    Force the number of nongauge parameters rather than rely on automated gauge optimization.
    
float_type : numpy dtype object, optional
    Numpy data type to use for floating point arrays.

{GATE_PENALTY_DESCRIPTION}

Returns
-------
CompositeScore
    The maximum score for `germs`, indicating how many parameters it amplifies.
""")
def compute_germ_set_score(germs, target_model=None, neighborhood=None,
                           neighborhood_size=5,
                           randomization_strength=1e-2, score_func='all',
                           op_penalty=0.0, l1_penalty=0.0, num_nongauge_params=None,
                           float_type=_np.cdouble, gate_penalty=None):
    pass

Method 2: add a dummy function that holds the documentation.

We can do something like

def GATE_PENALTY_DESCRIPTION():
    """
    gate_penalty : dict, optional (default None)
            An optional dictionary allowing the specification of gate-specific penalties to add for each instance
            of the specified gate(s) in each germ. Should be specified as a dictionary whose keys are strings
            corresponding to gate names, and whose values are the penalty factor to add for each instance of that
            gate. E.g. {'Gcnot':2} would correspond to a penalty term where each instance of a 'Gcnot' gate
            gets and additional 2 units added to the cost function for a candidate germ.
    """
    return

and then write a docstring more-or-less as usual

def compute_germ_set_score(germs, target_model=None, neighborhood=None,
                           neighborhood_size=5,
                           randomization_strength=1e-2, score_func='all',
                           op_penalty=0.0, l1_penalty=0.0, num_nongauge_params=None,
                           float_type=_np.cdouble, gate_penalty=None):
    """
    Calculate the score of a germ set with respect to a model.

    More precisely, this function computes the maximum score (roughly equal
    to the number of amplified parameters) for a cloud of models.
    If `target_model` is given, it serves as the center of the cloud,
    otherwise the cloud must be supplied directly via `neighborhood`.

    Parameters
    ----------
    germs : list
        The germ set

    target_model : Model, optional
        The target model, used to generate a neighborhood of randomized models.

    neighborhood : list of Models, optional
        The "cloud" of models for which scores are computed.  If not None, this
        overrides `target_model`, `neighborhood_size`, and `randomization_strength`.

    neighborhood_size : int, optional
        Number of randomized models to construct around `target_model`.

    randomization_strength : float, optional
        Strength of unitary randomizations, as passed to :meth:`target_model.randomize_with_unitary`.

    score_func : {'all', 'worst'}
        Sets the objective function for scoring the eigenvalues. If 'all',
        score is ``sum(1/input_array)``. If 'worst', score is ``1/min(input_array)``.

    op_penalty : float, optional
        Coefficient for a penalty linear in the sum of the germ lengths.

    l1_penalty : float, optional
        Coefficient for a penalty linear in the number of germs.

    num_nongauge_params : int, optional
        Force the number of nongauge parameters rather than rely on automated gauge optimization.
        
    float_type : numpy dtype object, optional
        Numpy data type to use for floating point arrays.

    gate_penalty
        Run `help(GATE_PENALTY_DESCRIPTION)` to see this parameter's description.

    Returns
    -------
    CompositeScore
        The maximum score for `germs`, indicating how many parameters it amplifies.
    """
    def score_fn(x): return _scoring.list_score(x, score_func=score_func)
    if neighborhood is None:
        neighborhood = [target_model.randomize_with_unitary(randomization_strength)
                        for n in range(neighborhood_size)]
    scores = [compute_composite_germ_set_score(score_fn, model=model,
                                               partial_germs_list=germs,
                                               op_penalty=op_penalty,
                                               l1_penalty=l1_penalty,
                                               num_nongauge_params=num_nongauge_params,
                                               float_type=float_type, 
                                               gate_penalty=gate_penalty,
                                               germ_list=germs)
              for model in neighborhood]

    return max(scores)

Method 3: docstring updates at end-of-file

We can define the module-wide constant GATE_PENALTY_DESCRIPTION like in method 1, and have the string literal "GATE_PENALTY_DESCRIPTION" in the docstring itself (also like method 1, but without the braces, quotes, or f-string). Then, at the end of the file we can do

compute_composite_germ_set_score.__doc__ = \
    compute_composite_germ_set_score.__doc__.replace('GATE_PENALTY_DESCRIPTION', GATE_PENALTY_DESCRIPTION)

The downside of this approach is that it adds a lot of messy code to the end of the file. If $N$ functions each have $M$ templated argument descriptions, then we'd have $MN$ calls to .replace(...).

[rileyjmurray]

This text appears three times in this file. Please break it out into a helper function.

[rileyjmurray]

The keyword arguments here are hard to read. I suggest something like the following.

class GermSelectionPenaltyTester(GermSelectionData, BaseCase):

    common_kwargs  : dict[str, Any] = dict(
        randomize=True, seed=1234, candidate_germ_counts={7:'all upto'},
        assume_real=True, float_type=np.double, mode='compactEVD', algorithm='greedy'
    )

    def setUp(self):
        super(GermSelectionData, self).setUp()

    def test_op_penalty_greedy(self):
        germs_no_penalty_cevd = germsel.find_germs(self.target_model, **self.common_kwargs)
        germs_penalty_cevd    = germsel.find_germs(self.target_model, **self.common_kwargs, algorithm_kwargs={'op_penalty':.1})        
        assert count_ops(germs_no_penalty_cevd) > count_ops(germs_penalty_cevd)


    def test_gate_penalty_greedy(self):
        kwargs = self.common_kwargs.copy()
        kwargs.pop('mode')
        germs_no_penalty_alljac = germsel.find_germs(self.target_model, **kwargs, mode='all-Jac')     
        germs_no_penalty_cevd   = germsel.find_germs(self.target_model, **kwargs, mode='compactEVD')
        germs_penalty_alljac = germsel.find_germs(
            self.target_model, **kwargs, mode='all-Jac',    algorithm_kwargs={'gate_penalty':{'Gxpi2':.1}}
        )        
        germs_penalty_cevd = germsel.find_germs(
            self.target_model, **kwargs, mode='compactEVD', algorithm_kwargs={'gate_penalty':{'Gxpi2':.1}}
        )
        assert count_gate(germs_no_penalty_alljac, 'Gxpi2') > count_gate(germs_penalty_alljac, 'Gxpi2')
        assert count_gate(germs_no_penalty_cevd,   'Gxpi2') > count_gate(germs_penalty_cevd,   'Gxpi2')
                                           
    def test_gate_penalty_grasp(self):
        kwargs = self.common_kwargs.copy()
        kwargs['mode']      = 'all-Jac'
        kwargs['algorithm'] = 'grasp'
        germs_gate_penalty_grasp = germsel.find_germs(
            self.target_model, **kwargs, algorithm_kwargs={'seed':1234, 'iterations':1, 'gate_penalty':{'Gxpi2':.2}}
        )
        germs_default_grasp = germsel.find_germs(
            self.target_model, **kwargs, algorithm_kwargs={'seed':1234, 'iterations':1}
        )
        assert count_gate(germs_gate_penalty_grasp, 'Gxpi2') < count_gate(germs_default_grasp, 'Gxpi2')

[rileyjmurray]

Missing blank line at EOF

[rileyjmurray]

@coreyostrove, is this is an appropriate PR to look into the complex-dtype issues with compactEVD germ selection #659? It's fine with me if you'd like to handle that separately, but I figured I'd point out the chance.