Source code for metaheuristic_designer.operators.masked_operator

"""
Operator that applies different operators to disjoint slices of the genotype.
"""

from __future__ import annotations
from typing import Iterable, Optional
from copy import copy
import numpy as np

from ..initializer import Initializer
from ..population import Population
from ..operator import Operator
from ..utils import MaskLike




class MaskedOperator(Operator):
    """Operator that partitions the genotype and applies different operators.

    A mask (integer array of length `vec_size`) specifies which operator
    (index into `op_list`) handles each gene.  This is used internally
    by :class:`ExtendedOperator` to separate the solution from auxiliary
    parameters.

    Parameters
    ----------
    op_list : list of Operator
        Operators to apply, one per mask index.
    mask : array of int
        Array of length `vec_size` assigning each gene to an operator.
    name : str, optional
        Display name; defaults to ``"Split (op_names)"``.
    **kwargs
        Forwarded to :class:`Operator`.
    """

    def __init__(self, op_list: Iterable[Operator], mask: MaskLike, name: str = None, **kwargs):
        if name is None:
            op_names = []
            for op in op_list:
                if not isinstance(op, Operator):
                    op_names.append("lambda_func")
                else:
                    op_names.append(op.name)

            joined_names = ", ".join(op_names)
            name = f"Split ({joined_names})"

        self.op_list = op_list
        super().__init__(name, mask=mask, **kwargs)



    def gather_params(self) -> dict:
        """Collect parameters from this operator and all sub-operators.

        Returns
        -------
        dict
            Flat dictionary with dotted keys.
        """

        all_params = self.get_params()
        for op in self.op_list:
            all_params.update(op.gather_params())

        return all_params




    def evolve(self, population: Population, initializer: Optional[Initializer] = None) -> Population:
        """Apply the appropriate operator to each slice of the genotype.

        Parameters
        ----------
        population : Population
            The current population.
        initializer : Initializer, optional
            The population initializer.

        Returns
        -------
        Population
            The modified population.
        """

        new_population = copy(population)

        # In masked_operator.py, inside the loop over op_list:
        for idx_op, op in enumerate(self.op_list):
            split_mask = self.params.mask == idx_op
            if np.any(split_mask):
                split_population = new_population.take_slice(split_mask)
                split_population = op.evolve(split_population, initializer)
                new_population = new_population.apply_slice(split_population, split_mask)

        return new_population




    def step(self, progress: float):
        """Update schedulable parameters and propagate to sub-operators.

        Parameters
        ----------
        progress : float
            Current progress of the algorithm (0-1).
        """

        super().step(progress)

        for op in self.op_list:
            if isinstance(op, Operator):
                op.step(progress)




    def get_state(self) -> dict:
        data = super().get_state()

        data["op_list"] = []

        for op in self.op_list:
            if isinstance(op, Operator):
                data["op_list"].append(op.get_state())
            else:
                data["op_list"].append("lambda_func")

        return data