text:
"""
The rule class is used for a specific application of a strategy on a tiling.
This is not something the user should implement, as it is just a wrapper for
calling the Strategy class and storing its results.

A CombinatorialSpecification is (more or less) a set of Rule.
"""
import abc
from functools import partial
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    Generic,
    Iterator,
    List,
    Optional,
    Sequence,
    Tuple,
    cast,
)

from sympy import Eq, Function

from ..combinatorial_class import CombinatorialClassType, CombinatorialObjectType
from ..exception import SanityCheckFailure, StrategyDoesNotApply
from .constructor import Constructor, DisjointUnion

if TYPE_CHECKING:
    from .strategy import (
        AbstractStrategy,
        Strategy,
        VerificationStrategy,
    )
    from .strategy_pack import StrategyPack

__all__ = ("Rule", "VerificationRule")

class AbstractRule(abc.ABC, Generic[CombinatorialClassType, CombinatorialObjectType]):
    """
    An instance of Rule is created by the __call__ method of strategy.
    """

    def __init__(
        self,
        strategy: "AbstractStrategy[CombinatorialClassType, CombinatorialObjectType]",
        comb_class: CombinatorialClassType,
        children: Optional[Tuple[CombinatorialClassType, ...]] = None,
    ):
        self.comb_class = comb_class
        self._strategy = strategy
        self.count_cache: Dict[Any, int] = {}
        self.obj_cache: Dict[Any, List[CombinatorialObjectType]] = {}
        self.subrecs: Optional[Tuple[Callable[..., int], ...]] = None
        self.subgenerators: Optional[
            Tuple[Callable[..., Iterator[CombinatorialObjectType]], ...]
        ] = (None)
        self.subsamplers: Optional[Tuple[Callable[..., CombinatorialObjectType], ...]]
        self._children = children
        self._non_empty_children: Optional[Tuple[CombinatorialClassType, ...]] = None

    @property
    def strategy(
        self,
    ) -> "AbstractStrategy[CombinatorialClassType, CombinatorialObjectType]":
        return self._strategy

    @property
    def ignore_parent(self) -> bool:
        """
        Return True if it is not worth expanding the parent/comb_class if
        this rule is found.
        """
        return self._strategy.ignore_parent

    @property
    def inferrable(self) -> bool:
        """
        Return True if the children could change using inferral strategies.
        """
        return self._strategy.inferrable

    @property
    def possibly_empty(self) -> bool:
        """
        Return True if it is possible that a child is empty.
        """
        return self._strategy.possibly_empty

    @property
    def workable(self) -> bool:
        """
        Return True if the children can expanded using other strategies.
        """
        return self._strategy.workable

    def set_subrecs(
        self,
        get_subrule: Callable[
            [CombinatorialClassType],
            "AbstractRule[CombinatorialClassType, CombinatorialObjectType]",
        ],
    ) -> None:
        """
        In general a rule comes in the context of a set of rules, e.g. in a
        CombinatorialSpecification. In order to count, generate, etc using the
        rule we must set the subfunctions that the children refer to.
        """
        self.subrecs = tuple(
            get_subrule(child).count_objects_of_size for child in self.children
        )
        self.subgenerators = tuple(
            get_subrule(child).generate_objects_of_size for child in self.children
        )
        self.subsamplers = tuple(
            get_subrule(child).random_sample_object_of_size for child in self.children
        )

    @property
    def children(self) -> Tuple[CombinatorialClassType, ...]:
        """
        Return the tuple of CombinatorialClass that are found by applying the
        decomposition function of the strategy to the parent.
        """
        if self._children is None:
            self._children = self._strategy.decomposition_function(self.comb_class)
            if self._children is None:
                raise StrategyDoesNotApply("{} does not apply".format(self._strategy))
        return self._children

    @property
    def formal_step(self) -> str:
        """
        Return a string with a short explanation about the rule.
        """
        return self._strategy.formal_step()

    @abc.abstractmethod
    def is_equivalence(self) -> bool:
        """
        Returns True if the rule is an equivalence.
        """

    def non_empty_children(self) -> Tuple[CombinatorialClassType, ...]:
        """
        Return the tuple of non-empty combinatorial classes that are found
        by applying the decomposition function.
        """
        if self._non_empty_children is None:
            self._non_empty_children = tuple(
                child for child in self.children if not child.is_empty()
            )
        return self._non_empty_children

    @abc.abstractmethod
    def count_objects_of_size(self, n: int, **parameters: int) -> int:
        """
        The function count the objects with respect to the parameters. The
        result is cached.
        """

    @abc.abstractmethod
    def get_equation(
        self, get_function: Callable[[CombinatorialClassType], Function]
    ) -> Eq:
        """
        Return the equation for the (ordinary) generating function.
        """

    @abc.abstractmethod
    def generate_objects_of_size(
        self, n: int, **parameters: int
    ) -> Iterator[CombinatorialObjectType]:
        """
        Generate the objects by using the underlying bijection between the
        parent and children.
        """

    @abc.abstractmethod
    def random_sample_object_of_size(
        self, n: int, **parameters: int
    ) -> CombinatorialObjectType:
        """Return a random objects of the give size."""

    def sanity_check(self, n: int, **parameters: int) -> bool:
        """
        Sanity check that this is a valid rule.

        Raise a SanityCheckFailure error if the sanity_check fails.
        """

        if isinstance(self, VerificationRule):
            # TODO: test more thoroughly
            return True

        def brute_force_count(
            comb_class: CombinatorialClassType, n: int, **parameters
        ) -> int:
            assert set(comb_class.extra_parameters) == set(
                parameters
            ), f"{comb_class.extra_parameters, set(parameters)}"
            return len(list(comb_class.objects_of_size(n, **parameters)))

        def brute_force_generation(
            comb_class: CombinatorialClassType, n: int, **parameters: int
        ) -> Iterator[CombinatorialObjectType]:
            assert set(comb_class.extra_parameters) == set(
                parameters
            ), f"{comb_class.extra_parameters, set(parameters)}"
            yield from comb_class.objects_of_size(n, **parameters)

        actual_count = brute_force_count(self.comb_class, n, **parameters)
        temprec = self.subrecs
        self.subrecs = tuple(
            partial(brute_force_count, child) for child in self.children
        )
        rule_count = self.count_objects_of_size(n, **parameters)
        self.subrecs = temprec
        params_str = ", ".join(
            [f"n = {n}"] + [f"{p} = {v}" for p, v in parameters.items()]
        )
        if actual_count != rule_count:
            raise SanityCheckFailure(
                f"The following rule failed sanity check:\n"
                f"{self}\n"
                f"Failed with parameters:\n"
                f"{params_str}\n"
                f"The actual count is {actual_count}.\n"
                f"The rule count is {rule_count}.",
            )
        if any(
            comb_class.extra_parameters
            for comb_class in [self.comb_class, *self.children]
        ):
            # Can't sanity check generation of objects for classes with extra
            # TODO test more thoroughly
            return True
        tempgen = self.subgenerators
        self.subgenerators = tuple(
            partial(brute_force_generation, child) for child in self.children
        )
        try:
            rule_objects = set(list(self.generate_objects_of_size(n, **parameters)))
        except NotImplementedError:
            # Skipping testing rules that have not implemented object generation.
            return True
        self.subgenerators = tempgen
        actual_objects = set(
            list(brute_force_generation(self.comb_class, n, **parameters))
        )
        if actual_objects != rule_objects:
            raise SanityCheckFailure(
                f"The following rule failed sanity check:\n"
                f"{self}\n"
                f"Failed with parameters:\n"
                f"{params_str}\n"
                f"The rule generated the wrong objects."
            )
        return True

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, AbstractRule):
            return NotImplemented
        return (
            isinstance(other, self.__class__)
            and self.comb_class == other.comb_class
            and self._strategy == other._strategy
        )

    def __str__(self) -> str:
        def frontpad(res, height):
            n = max(len(s) for s in res)
            for _ in range(height - len(res)):
                res.append(" " * n)

        def backpad(res):
            n = max(len(s) for s in res)
            for i, s in enumerate(res):
                res[i] = s + " " * (n - len(s))

        def join(res, new):
            backpad(new)
            backpad(res)
            frontpad(res, len(new))
            for i, s in enumerate(new):
                res[i] += s

        res = str(self.comb_class).split("\n")
        backpad(res)
        if isinstance(self, Rule):
            children = [str(child).split("\n") for child in self.children]
            symbol_height = min(1, len(res) - 1)
            eq_symbol = (
                ["     " for i in range(symbol_height)]
                + ["  {}  ".format(self.strategy.get_eq_symbol())]
                + ["     " for i in range(symbol_height)]
            )
            join(res, eq_symbol)
            join(res, children[0])
            if len(children) > 1:
                op_symbol = (
                    ["     " for i in range(symbol_height)]
                    + ["  {}  ".format(self.strategy.get_op_symbol())]
                    + ["     " for i in range(symbol_height)]
                )
                for child in children[1:]:
                    join(res, op_symbol)
                    join(res, child)
        return f"{self.formal_step}\n" + "\n".join(x for x in res)

class Rule(AbstractRule[CombinatorialClassType, CombinatorialObjectType]):
    def __init__(
        self,
        strategy: "Strategy[CombinatorialClassType, CombinatorialObjectType]",
        comb_class: CombinatorialClassType,
        children: Optional[Tuple[CombinatorialClassType, ...]] = None,
    ):
        super().__init__(strategy, comb_class, children)
        self._constructor: Optional[Constructor] = None

    @property
    def strategy(self) -> "Strategy[CombinatorialClassType, CombinatorialObjectType]":
        return cast(
            "Strategy[CombinatorialClassType, CombinatorialObjectType]",
            super().strategy,
        )

    @property
    def constructor(
        self,
    ) -> Constructor[CombinatorialClassType, CombinatorialObjectType]:
        """
        Return the constructor, that contains all the information about how to
        count/generate objects from the rule.
        """
        if self._constructor is None:
            self._constructor = self.strategy.constructor(
                self.comb_class, self.children
            )
            if self._constructor is None:
                raise StrategyDoesNotApply("{} does not apply".format(self.strategy))
        return self._constructor

    def is_equivalence(self):
        return self.strategy.can_be_equivalent() and len(self.non_empty_children()) == 1

    def backward_map(
        self, objs: Tuple[Optional[CombinatorialObjectType], ...]
    ) -> CombinatorialObjectType:
        """
        This encodes the backward map of the underlying bijection that the
        strategy implies.
        """
        return self.strategy.backward_map(self.comb_class, objs, self.children)

    def forward_map(
        self, obj: CombinatorialObjectType
    ) -> Tuple[Optional[CombinatorialObjectType], ...]:
        """
        This encodes the forward map of the underlying bijection that the
        strategy implies.
        """
        return self.strategy.forward_map(self.comb_class, obj, self.children)

    def to_equivalence_rule(self) -> "EquivalenceRule":
        """
        Return the reverse rule. At this stage, reverse rules can only be
        created for equivalence rules.
        """
        assert (
            self.is_equivalence()
        ), f"EquivalenceRule can only be created for equivalence rules\n{self}"
        return EquivalenceRule(self)

    def to_reverse_rule(self) -> "ReverseRule":
        """
        Return the reverse rule. At this stage, reverse rules can only be
        created for equivalence rules.
        """
        assert (
            self.is_equivalence()
        ), "reverse rule can only be created for equivalence rules"
        return ReverseRule(self)

    def count_objects_of_size(self, n: int, **parameters: int) -> int:
        """
        The function count the objects with respect to the parameters. The
        result is cached.
        """
        key = (n,) + tuple(sorted(parameters.items()))
        assert sorted(["n"] + list(parameters)) == sorted(
            ("n",) + self.comb_class.extra_parameters
        ), (
            "parameters did not match comb_class for the rule \n{}\nparameters "
            "given: {}\n comb_class_parameters: {}".format(
                self, list(parameters), self.comb_class.extra_parameters,
            )
        )
        res = self.count_cache.get(key)
        if res is None:
            assert (
                self.subrecs is not None
            ), "you must call the set_subrecs function first"
            try:
                res = self.constructor.get_recurrence(self.subrecs, n, **parameters)
            except AssertionError:
                raise ValueError(f"issue with rule:\n {self}")
            self.count_cache[key] = res
            # THE FOLLOWING CODE SNIPPET IS FOR DEBUGGING PURPOSES
        #     if self.comb_class.extra_parameters:
        #         print(self)
        #         print("n =", n, parameters)
        #         if hasattr(self.constructor, "extra_parameters"):
        #             print("parent->child params:", self.constructor.extra_parameters)
        #         if hasattr(self.constructor, "split_parameters"):
        #             print("parent->child params:", self.constructor.split_parameters)
        #         if hasattr(self.constructor, "zeroes"):
        #             print("zeroes:", self.constructor.zeroes)
        #         if hasattr(self.constructor, "get_extra_parameters"):
        #             print(
        #                 "parameters_passed:",
        #                 self.constructor.get_extra_parameters(n, **parameters),
        #             )
        #         fusion_attrs = [
        #             "extra_parameters",
        #             "reversed_extra_parameters",
        #             "fuse_parameter",
        #             "left_sided_parameters",
        #             "right_sided_parameters",
        #             "both_sided_parameters",
        #             "rec_function",
        #             "fusion_type",
        #             "predeterminable_left_right_points",
        #         ]
        #         for string in fusion_attrs:
        #             if hasattr(self.constructor, string):
        #                 print(string + ":", getattr(self.constructor, string))
        #         if isinstance(self.constructor, CartesianProduct):
        #             print(list(self.constructor._valid_compositions(n, **parameters)))
        #     print("result:", res)
        # assert res == len(list(self.comb_class.objects_of_size(n, **parameters))), (
        #     "counting failed for the rule \n{}\nparameters: n = {}, {}\n"
        #     "computed {}, actual {}".format(
        #         self,
        #         n,
        #         parameters,
        #         res,
        #         len(list(self.comb_class.objects_of_size(n, **parameters))),
        #     )
        # )
        return res

    def get_equation(
        self, get_function: Callable[[CombinatorialClassType], Function]
    ) -> Eq:
        """
        Return the equation for the (ordinary) generating function.
        """
        lhs_func = get_function(self.comb_class)
        rhs_funcs = tuple(get_function(comb_class) for comb_class in self.children)
        return self.constructor.get_equation(lhs_func, rhs_funcs)

    def generate_objects_of_size(
        self, n: int, **parameters: int
    ) -> Iterator[CombinatorialObjectType]:
        """
        Generate the objects by using the underlying bijection between the
        parent and children.
        """
        key = (n,) + tuple(sorted(parameters.items()))
        res = self.obj_cache.get(key)
        if res is not None:
            yield from res
            return
        assert (
            self.subgenerators is not None
        ), "you must call the set_subrecs function first"
        res = []
        for subobjs in self.constructor.get_sub_objects(
            self.subgenerators, n, **parameters
        ):
            obj = self.backward_map(subobjs)
            yield obj
            res.append(obj)
        self.obj_cache[key] = res

    def random_sample_object_of_size(
        self, n: int, **parameters: int
    ) -> CombinatorialObjectType:
        assert self.subrecs is not None, "you must call the set_subrecs function first"
        assert (
            self.subsamplers is not None
        ), "you must call the set_subrecs function first"
        total_count = self.count_objects_of_size(n=n, **parameters)
        subobjs = self.constructor.random_sample_sub_objects(
            total_count, self.subsamplers, self.subrecs, n, **parameters
        )
        obj = self.backward_map(subobjs)
        return obj

class EquivalenceRule(Rule[CombinatorialClassType, CombinatorialObjectType]):
    def __init__(self, rule: Rule):
        non_empty_children = rule.non_empty_children()
        assert rule.is_equivalence(), "not an equivalence rule: {}".format(str(rule))
        child = non_empty_children[0]
        super().__init__(rule.strategy, rule.comb_class, (child,))
        self.child_idx = rule.children.index(child)
        self.actual_children = rule.children
        self._constructor: Optional[DisjointUnion] = None

    @property
    def constructor(self,) -> DisjointUnion:
        """
        Return the constructor, that contains all the information about how to
        count/generate objects from the rule.
        """
        if self._constructor is None:
            original_constructor = cast(
                DisjointUnion,
                self.strategy.constructor(self.comb_class, self.actual_children),
            )
            assert isinstance(original_constructor, DisjointUnion)
            self._constructor = DisjointUnion(
                self.comb_class,
                self.children,
                (original_constructor.extra_parameters[self.child_idx],),
            )
        return self._constructor

    @staticmethod
    def is_equivalence() -> bool:
        return True

    @property
    def formal_step(self) -> str:
        return "{} but only the child at index {} is non-empty".format(
            self.strategy.formal_step(), self.child_idx
        )

    def backward_map(
        self, objs: Tuple[Optional[CombinatorialObjectType], ...]
    ) -> CombinatorialObjectType:
        actual_objs = tuple(
            objs[0] if i == self.child_idx else None
            for i in range(len(self.actual_children))
        )
        return self.strategy.backward_map(
            self.comb_class, actual_objs, self.actual_children
        )

    def forward_map(
        self, obj: CombinatorialObjectType
    ) -> Tuple[Optional[CombinatorialObjectType], ...]:
        return (
            self.strategy.forward_map(self.comb_class, obj, self.actual_children)[
                self.child_idx
            ],
        )

class EquivalencePathRule(Rule[CombinatorialClassType, CombinatorialObjectType]):
    """
    A class for shortening a chain of equivalence rules into a single Rule.
    """

    def __init__(self, rules: Sequence[Rule]):
        assert all(rule.is_equivalence() for rule in rules)
        super().__init__(rules[0].strategy, rules[0].comb_class, rules[-1].children)
        self.rules = rules
        self._constructor: Optional[DisjointUnion] = None

    @property
    def constructor(self) -> DisjointUnion:
        if self._constructor is None:
            extra_parameters: Dict[str, str] = {
                k: k for k in self.comb_class.extra_parameters
            }
            for rule in self.rules:
                original_constructor = rule.constructor
                assert isinstance(original_constructor, DisjointUnion)
                rules_parameters = original_constructor.extra_parameters[0]
                extra_parameters = {
                    parent_var: rules_parameters[child_var]
                    for parent_var, child_var in extra_parameters.items()
                    if child_var in rules_parameters
                }
            fixed_values = {
                k: 0
                for k in self.children[0].extra_parameters
                if k not in extra_parameters.values()
            }
            self._constructor = DisjointUnion(
                self.comb_class, self.children, (extra_parameters,), (fixed_values,)
            )
        return self._constructor

    @staticmethod
    def is_equivalence() -> bool:
        return True

    @property
    def formal_step(self) -> str:
        return ", then ".join(rule.formal_step for rule in self.rules)

    def eqv_path_rules(self) -> List[Tuple[CombinatorialClassType, Rule]]:
        """
        Returns the list of (parent, rule) pairs that make up the equivalence
        path.
        """
        eqv_path_rules = []
        curr = self.comb_class
        for rule in self.rules:
            eqv_path_rules.append((curr, rule))
            curr = rule.children[0]
        return eqv_path_rules

    def backward_map(
        self, objs: Tuple[Optional[CombinatorialObjectType], ...]
    ) -> CombinatorialObjectType:
        res = cast(Tuple[CombinatorialObjectType], objs)
        for rule in reversed(self.rules):
            res = (rule.backward_map(res),)
        return res[0]

    def forward_map(
        self, obj: CombinatorialObjectType
    ) -> Tuple[Optional[CombinatorialObjectType], ...]:
        res = cast(CombinatorialObjectType, obj)
        for rule in reversed(self.rules):
            res = cast(CombinatorialObjectType, rule.forward_map(res)[0])
        return (res,)

    def __eq__(self, other) -> bool:
        return super().__eq__(other) and [r.strategy for r in self.rules] == [
            r.strategy for r in other.rules
        ]

    def __str__(self) -> str:
        def frontpad(res, height):
            n = max(len(s) for s in res)
            for _ in range(height - len(res)):
                res.append(" " * n)

        def backpad(res):
            n = max(len(s) for s in res)
            for i, s in enumerate(res):
                res[i] = s + " " * (n - len(s))

        def join(res, new):
            backpad(new)
            backpad(res)
            frontpad(res, len(new))
            for i, s in enumerate(new):
                res[i] += s

        res = str(self.comb_class).split("\n")
        backpad(res)
        comb_classes = [str(rule.children[0]).split("\n") for rule in self.rules]
        symbol_height = min(1, len(res) - 1)
        eq_symbol = (
            ["     " for i in range(symbol_height)]
            + ["  {}  ".format("=")]
            + ["     " for i in range(symbol_height)]
        )
        for comb_class in comb_classes:
            join(res, eq_symbol)
            join(res, comb_class)
        return "{}\n".format(self.formal_step) + "\n".join(x for x in res)

class ReverseRule(Rule[CombinatorialClassType, CombinatorialObjectType]):
    """
    A class for creating the reverse of an equivalence rule.
    """

    def __init__(self, rule: Rule):
        assert (
            rule.is_equivalence()
        ), "reversing a rule only works for equivalence rules"
        super().__init__(rule.strategy, rule.children[0], (rule.comb_class,))
        self.original_rule = rule
        self._constructor: Optional[DisjointUnion] = None

    @property
    def constructor(self) -> DisjointUnion:
        if self._constructor is None:
            original_constructor = self.original_rule.constructor
            assert isinstance(original_constructor, DisjointUnion), (
                "reverse rule coming from non disjoint union strategy - "
                "you'll need to update the ReverseRule constructor!"
            )
            flipped_extra_params = {
                b: a for a, b in original_constructor.extra_parameters[0].items()
            }
            fixed_values = {
                k: 0
                for k in self.children[0].extra_parameters
                if k not in flipped_extra_params.values()
            }
            self._constructor = DisjointUnion(
                self.comb_class, self.children, (flipped_extra_params,), (fixed_values,)
            )
        return self._constructor

    @staticmethod
    def is_equivalence() -> bool:
        return True

    @property
    def formal_step(self) -> str:
        return "reverse of '{}'".format(self.strategy.formal_step())

    def backward_map(
        self, objs: Tuple[Optional[CombinatorialObjectType], ...]
    ) -> CombinatorialObjectType:
        return cast(CombinatorialObjectType, self.original_rule.forward_map(objs[0])[0])

    def forward_map(
        self, obj: CombinatorialObjectType
    ) -> Tuple[Optional[CombinatorialObjectType], ...]:
        return (self.original_rule.backward_map((obj,)),)

class VerificationRule(AbstractRule[CombinatorialClassType, CombinatorialObjectType]):
    """
    A class for verification rules. In particular, the children will be an
    empty tuple if it applies, else None.
    """

    def __init__(
        self,
        strat: "VerificationStrategy[CombinatorialClassType,CombinatorialObjectType]",
        comb_class: CombinatorialClassType,
        children: Optional[Tuple[CombinatorialClassType, ...]] = None,
    ):
        assert not children
        super().__init__(strat, comb_class, children)

    @property
    def strategy(
        self,
    ) -> "VerificationStrategy[CombinatorialClassType, CombinatorialObjectType]":
        return cast(
            "VerificationStrategy[CombinatorialClassType, CombinatorialObjectType]",
            super().strategy,
        )

    @staticmethod
    def is_equivalence() -> bool:
        return False

    def count_objects_of_size(self, n: int, **parameters: int) -> int:
        key = (n,) + tuple(sorted(parameters.items()))
        assert set(parameters) == set(self.comb_class.extra_parameters)
        res = self.count_cache.get(key)
        if res is None:
            res = self.strategy.count_objects_of_size(self.comb_class, n, **parameters)
            self.count_cache[key] = res
        # # THE FOLLOWING CODE SNIPPET IS FOR DEBUGGING PURPOSES
        # assert res == len(list(self.comb_class.objects_of_size(n, **parameters))), (
        #     "counting failed for the rule \n{}\nparameters: n = {}, {}\n"
        #     "computed {}, actual {}".format(
        #         self,
        #         n,
        #         parameters,
        #         res,
        #         len(list(self.comb_class.objects_of_size(n, **parameters))),
        #     )
        # )
        return res

    def pack(self) -> "StrategyPack":
        return self.strategy.pack(self.comb_class)

    def get_equation(
        self,
        get_function: Callable[[CombinatorialClassType], Function],
        funcs: Optional[Dict[CombinatorialClassType, Function]] = None,
    ) -> Eq:
        lhs_func = get_function(self.comb_class)
        return Eq(lhs_func, self.strategy.get_genf(self.comb_class, funcs))

    def generate_objects_of_size(
        self, n: int, **parameters: int
    ) -> Iterator[CombinatorialObjectType]:
        key = (n,) + tuple(sorted(parameters.items()))
        res = self.obj_cache.get(key)
        if res is not None:
            yield from res
            return
        res = []
        for obj in self.strategy.generate_objects_of_size(
            self.comb_class, n, **parameters
        ):
            yield obj
            res.append(obj)
        self.obj_cache[key] = res

    def random_sample_object_of_size(
        self, n: int, **parameters: int
    ) -> CombinatorialObjectType:
        return self.strategy.random_sample_object_of_size(
            self.comb_class, n, **parameters
        )

reason: from sympy


text:
#!/usr/bin/env python

# Copyright 2016-2019 Biomedical Imaging Group Rotterdam, Departments of
# Medical Informatics and Radiology, Erasmus MC, Rotterdam, The Netherlands
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
from sklearn.base import BaseEstimator, ClassifierMixin
from sklearn.utils.validation import check_is_fitted
from sklearn.utils.multiclass import unique_labels
import WORC.classification.RankedSVM as RSVM

class RankedSVM(BaseEstimator, ClassifierMixin):
    """ An example classifier which implements a 1-NN algorithm.

    Parameters
    ----------
    demo_param : str, optional
        A parameter used for demonstation of how to pass and store paramters.

    Attributes
    ----------
    X_ : array, shape = [n_samples, n_features]
        The input passed during :meth:`fit`
    y_ : array, shape = [n_samples]
        The labels passed during :meth:`fit`
    """
    def __init__(self, cost=1, lambda_tol=1e-6,
                 norm_tol=1e-4, max_iter=500, svm='Poly', gamma=0.05,
                 coefficient=0.05, degree=3):
        self.cost = cost
        self.lambda_tol = lambda_tol
        self.norm_tol = norm_tol
        self.max_iter = max_iter
        self.svm = svm
        self.gamma = gamma
        self.coefficient = coefficient
        self.degree = 3

    def fit(self, X, y):
        """A reference implementation of a fitting function for a classifier.

        Parameters
        ----------
        X : array-like, shape = [n_samples, n_features]
            The training input samples.
        y : array-like, shape = [n_samples]
            The target values. An array of int.

        Returns
        -------
        self : object
            Returns self.
        """

        # Store the classes seen during fit
        self.classes_ = unique_labels(y)

        # RankedSVM requires a very specific format of y
        # Each row should represent a label, consisiting of ones and minus ones
        y = np.transpose(y).astype(np.int16)
        y[y == 0] = -1
        self.X_ = X
        self.y_ = y
        self.num_class = y.shape[0]

        Weights, Bias, SVs =\
            RSVM.RankSVM_train(train_data=X,
                               train_target=y,
                               cost=self.cost,
                               lambda_tol=self.lambda_tol,
                               norm_tol=self.norm_tol,
                               max_iter=self.max_iter,
                               svm=self.svm, gamma=self.gamma,
                               coefficient=self.coefficient,
                               degree=self.degree)

        self.Weights = Weights
        self.Bias = Bias
        self.SVs = SVs

        return self

    def predict(self, X, y=None):
        """ A reference implementation of a prediction for a classifier.

        Parameters
        ----------
        X : array-like of shape = [n_samples, n_features]
            The input samples.

        Returns
        -------
        y : array of int of shape = [n_samples]
            The label for each sample is the label of the closest sample
            seen udring fit.
        """
        # Check is fit had been called
        check_is_fitted(self, ['X_', 'y_'])

        _, Predicted_Labels =\
            RSVM.RankSVM_test(test_data=X,
                              num_class=self.num_class,
                              Weights=self.Weights,
                              Bias=self.Bias,
                              SVs=self.SVs,
                              svm=self.svm, gamma=self.gamma,
                              coefficient=self.coefficient,
                              degree=self.degree)

        return Predicted_Labels

    def predict_proba(self, X, y):
        """ A reference implementation of a prediction for a classifier.

        Parameters
        ----------
        X : array-like of shape = [n_samples, n_features]
            The input samples.

        Returns
        -------
        y : array of int of shape = [n_samples]
            The label for each sample is the label of the closest sample
            seen udring fit.
        """
        # Check is fit had been called
        check_is_fitted(self, ['X_', 'y_'])

        Probs, _ =\
            RSVM.RankSVM_test(test_data=X,
                              num_class=self.num_class,
                              Weights=self.Weights,
                              Bias=self.Bias,
                              svm=self.svm, gamma=self.gamma,
                              coefficient=self.coefficient,
                              degree=self.degree)

        return Probs

ext: py

max_stars_repo_name: Sikerdebaard/PREDICTFastr

max_issues_repo_name: Sikerdebaard/PREDICTFastr

max_forks_repo_name: Sikerdebaard/PREDICTFastr

reason: import numpy


text:
import numpy as np
from scratch_ml.utils import euclidean_distance

class KMeans():
    """A simple clustering method that forms k clusters by iteratively reassigning
    samples to the closest centroids and after that moves the centroids to the center
    of the new formed clusters."""

    def __init__(self, k=2, iterations=500):
        self.k = k
        self.iterations = iterations

    def _init_random_centroids(self, x):
        """Initialize the centroids"""
        n_samples, n_features = np.shape(x)
        centroids = np.zeros((self.k, n_features))
        for i in range(self.k):
            centroid = x[np.random.choice(range(n_samples))]
            centroids[i] = centroid
        return centroids

    def _closest_centroid(self, sample, centroids):
        """Return the index of the closest centroid to the sample"""
        closest_idx = 0
        closest_dist = float('inf')
        for i, centroid in enumerate(centroids):
            distance = euclidean_distance(sample, centroid)
            if distance < closest_dist:
                closest_idx = i
                closest_dist = distance
        return closest_idx

    def _create_clusters(self, centroids, x):
        """Assign the samples to the closest centroids to create clusters"""
        clusters = [[] for _ in range(self.k)]
        for sample_idx, sample in enumerate(x):
            centroid_idx = self._closest_centroid(sample, centroids)
            clusters[centroid_idx].append(sample_idx)
        return clusters

    def _calculate_centroids(self, clusters, x):
        """Calculate new centroids"""
        n_features = np.shape(x)[1]
        centroids = np.zeros((self.k, n_features))
        for i, cluster in enumerate(clusters):
            centroid = np.mean(x[cluster], axis=0)
            centroids[i] = centroid
        return centroids

    def _get_cluster_labels(self, clusters, x):
        """Classify samples as the index of their clusters"""
        y_pred = np.zeros(np.shape(x)[0])
        for cluster_i, cluster in enumerate(clusters):
            for sample_i in cluster:
                y_pred[sample_i] = cluster_i
        return y_pred

    def predict(self, x):
        centroids = self._init_random_centroids(x)
        for i in range(self.iterations):
            clusters = self._create_clusters(centroids, x)
            prev_centroids = centroids
            centroids = self._calculate_centroids(clusters, x)
            diff = centroids - prev_centroids
            if not diff.any():
                break
        return self._get_cluster_labels(clusters, x)


text:
"""
=========================================
Nested versus non-nested cross-validation
=========================================

This example compares non-nested and nested cross-validation strategies on a
classifier of the iris data set. Nested cross-validation (CV) is often used to
train a model in which hyperparameters also need to be optimized. Nested CV
estimates the generalization error of the underlying model and its
(hyper)parameter search. Choosing the parameters that maximize non-nested CV
biases the model to the dataset, yielding an overly-optimistic score.

Model selection without nested CV uses the same data to tune model parameters
and evaluate model performance. Information may thus "leak" into the model
and overfit the data. The magnitude of this effect is primarily dependent on
the size of the dataset and the stability of the model. See Cawley and Talbot
[1]_ for an analysis of these issues.

To avoid this problem, nested CV effectively uses a series of
train/validation/test set splits. In the inner loop (here executed by
:class:`GridSearchCV <sklearn.model_selection.GridSearchCV>`), the score is
approximately maximized by fitting a model to each training set, and then
directly maximized in selecting (hyper)parameters over the validation set. In
the outer loop (here in :func:`cross_val_score
<sklearn.model_selection.cross_val_score>`), generalization error is estimated
by averaging test set scores over several dataset splits.

The example below uses a support vector classifier with a non-linear kernel to
build a model with optimized hyperparameters by grid search. We compare the
performance of non-nested and nested CV strategies by taking the difference
between their scores.

.. topic:: See Also:

    - :ref:`cross_validation`
    - :ref:`grid_search`

.. topic:: References:

    .. [1] `Cawley, G.C.; Talbot, N.L.C. On over-fitting in model selection and
     subsequent selection bias in performance evaluation.
     J. Mach. Learn. Res 2010,11, 2079-2107.
     <http://jmlr.csail.mit.edu/papers/volume11/cawley10a/cawley10a.pdf>`_

"""
from sklearn.datasets import load_iris
from matplotlib import pyplot as plt
from sklearn.svm import SVC
from sklearn.model_selection import GridSearchCV, cross_val_score, KFold
import numpy as np

print(__doc__)

# Number of random trials
NUM_TRIALS = 30

# Load the dataset
iris = load_iris()
X_iris = iris.data
y_iris = iris.target

# Set up possible values of parameters to optimize over
p_grid = {"C": [1, 10, 100],
          "gamma": [.01, .1]}

# We will use a Support Vector Classifier with "rbf" kernel
svm = SVC(kernel="rbf")

# Arrays to store scores
non_nested_scores = np.zeros(NUM_TRIALS)
nested_scores = np.zeros(NUM_TRIALS)

# Loop for each trial
for i in range(NUM_TRIALS):

    # Choose cross-validation techniques for the inner and outer loops,
    # independently of the dataset.
    # E.g "LabelKFold", "LeaveOneOut", "LeaveOneLabelOut", etc.
    inner_cv = KFold(n_splits=4, shuffle=True, random_state=i)
    outer_cv = KFold(n_splits=4, shuffle=True, random_state=i)

    # Non_nested parameter search and scoring
    clf = GridSearchCV(estimator=svm, param_grid=p_grid, cv=inner_cv)
    clf.fit(X_iris, y_iris)
    non_nested_scores[i] = clf.best_score_

    # Nested CV with parameter optimization
    nested_score = cross_val_score(clf, X=X_iris, y=y_iris, cv=outer_cv)
    nested_scores[i] = nested_score.mean()

score_difference = non_nested_scores - nested_scores

print("Average difference of {0:6f} with std. dev. of {1:6f}."
      .format(score_difference.mean(), score_difference.std()))

# Plot scores on each trial for nested and non-nested CV
plt.figure()
plt.subplot(211)
non_nested_scores_line, = plt.plot(non_nested_scores, color='r')
nested_line, = plt.plot(nested_scores, color='b')
plt.ylabel("score", fontsize="14")
plt.legend([non_nested_scores_line, nested_line],
           ["Non-Nested CV", "Nested CV"],
           bbox_to_anchor=(0, .4, .5, 0))
plt.title("Non-Nested and Nested Cross Validation on Iris Dataset",
          x=.5, y=1.1, fontsize="15")

# Plot bar chart of the difference.
plt.subplot(212)
difference_plot = plt.bar(range(NUM_TRIALS), score_difference)
plt.xlabel("Individual Trial #")
plt.legend([difference_plot],
           ["Non-Nested CV - Nested CV Score"],
           bbox_to_anchor=(0, 1, .8, 0))
plt.ylabel("score difference", fontsize="14")

plt.show()

max_issues_repo_name: amirothman/scikit-learn-course

max_forks_repo_name: amirothman/scikit-learn-course


text:
r"""
Quantum gate synthesis
======================

This demonstration works through the process used to produce the gate
synthesis results presented in `"Machine learning method for state
preparation and gate synthesis on photonic quantum
computers" <https://arxiv.org/abs/1807.10781>`__.

This tutorial uses the TensorFlow backend of Strawberry Fields, giving us
access to a number of additional functionalities including: GPU integration,
automatic gradient computation, built-in optimization algorithms, and other
machine learning tools.

Variational quantum circuits
----------------------------

A key element of machine learning is optimization. We can use
TensorFlow's automatic differentiation tools to optimize the parameters
of variational quantum circuits constructed using Strawberry Fields. In
this approach, we fix a circuit architecture where the states, gates,
and/or measurements may have learnable parameters :math:`\vec{\theta}`
associated with them. We then define a loss function based on the output
state of this circuit. In this case, we define a loss function such that the
action of the variational quantum circuit is close to some specified target
unitary.

.. note::

    For more details on the TensorFlow backend in Strawberry Fields, please see
    :ref:`machine_learning_tutorial`.

For arbitrary gate synthesis using optimization, we need to make use of
a quantum circuit with a layer structure that is **universal** - that
is, by 'stacking' the layers, we can guarantee that we can produce *any*
CV state with at-most polynomial overhead. Therefore, the architecture
we choose must consist of layers with each layer containing
parameterized Gaussian *and* non-Gaussian gates. **The non-Gaussian
gates provide both the nonlinearity and the universality of the model.**
To this end, we employ the CV quantum neural network architecture
described below:

.. figure:: https://i.imgur.com/NEsaVIX.png
   :alt: layer

Here,

-  :math:`\mathcal{U}_i(\theta_i,\phi_i)` is an N-mode linear optical
   interferometer composed of two-mode beamsplitters
   :math:`BS(\theta,\phi)` and single-mode rotation gates
   :math:`R(\phi)=e^{i\phi\hat{n}}`,

-  :math:`\mathcal{D}(\alpha_i)` are single mode displacements in the
   phase space by complex value :math:`\alpha_i`,

-  :math:`\mathcal{S}(r_i, \phi_i)` are single mode squeezing operations
   of magnitude :math:`r_i` and phase :math:`\phi_i`, and

-  :math:`\Phi(\lambda_i)` is a single mode non-Gaussian operation, in
   this case chosen to be the Kerr interaction
   :math:`\mathcal{K}(\kappa_i)=e^{i\kappa_i\hat{n}^2}` of strength
   :math:`\kappa_i`.

Hyperparameters
---------------

First, we must define the **hyperparameters** of our layer structure:

-  ``cutoff``: the simulation Fock space truncation we will use in the
   optimization. The TensorFlow backend will perform numerical
   operations in this truncated Fock space when performing the
   optimization.

-  ``depth``: The number of layers in our variational quantum
   circuit. As a general rule, increasing the number of layers (and
   thus, the number of parameters we are optimizing over) increases the
   optimizer's chance of finding a reasonable local minimum in the
   optimization landscape.

-  ``reps``: the number of steps in the optimization routine performing
   gradient descent

Some other optional hyperparameters include:

-  The standard deviation of initial parameters. Note that we make a
   distinction between the standard deviation of *passive* parameters
   (those that preserve photon number when changed, such as phase
   parameters), and *active* parameters (those that introduce or remove
   energy from the system when changed).
"""
import numpy as np

import strawberryfields as sf
from strawberryfields.ops import *
from strawberryfields.utils import operation

# Cutoff dimension
cutoff = 10

# gate cutoff
gate_cutoff = 4

# Number of layers
depth = 15

# Number of steps in optimization routine performing gradient descent
reps = 200

# Learning rate
lr = 0.025

# Standard deviation of initial parameters
passive_sd = 0.1
active_sd = 0.001

######################################################################
# Note that, unlike in state learning, we must also specify a *gate
# cutoff* :math:`d`. This restricts the target unitary to its action on a
# :math:`d`-dimensional subspace of the truncated Fock space, where
# :math:`d\leq D`, where :math:`D` is the overall simulation Fock basis
# cutoff. As a result, we restrict the gate synthesis optimization to only
# :math:`d` input-output relations.
#
# The layer parameters :math:`\vec{\theta}`
# -----------------------------------------
#
# We use TensorFlow to create the variables corresponding to the gate
# parameters. Note that each variable has shape ``(depth,)``, with each
# individual element representing the gate parameter in layer :math:`i`.
#

import tensorflow as tf

# set the random seed
tf.random.set_seed(42)
np.random.seed(42)

# squeeze gate
sq_r = tf.random.normal(shape=[depth], stddev=active_sd)
sq_phi = tf.random.normal(shape=[depth], stddev=passive_sd)

# displacement gate
d_r = tf.random.normal(shape=[depth], stddev=active_sd)
d_phi = tf.random.normal(shape=[depth], stddev=passive_sd)

# rotation gates
r1 = tf.random.normal(shape=[depth], stddev=passive_sd)
r2 = tf.random.normal(shape=[depth], stddev=passive_sd)

# kerr gate
kappa = tf.random.normal(shape=[depth], stddev=active_sd)

######################################################################
# For convenience, we store the TensorFlow variables representing the weights as
# a tensor:

weights = tf.convert_to_tensor([r1, sq_r, sq_phi, r2, d_r, d_phi, kappa])
weights = tf.Variable(tf.transpose(weights))

######################################################################
# Since we have a depth of 15 (so 15 layers), and each layer takes
# 7 different types of parameters, the final shape of our weights
# array should be :math:`\text{depth}\times 7` or ``(15, 7)``:

print(weights.shape)

######################################################################
# .. rst-class:: sphx-glr-script-out
#
#  Out:
#
#  .. code-block:: none
#
#     (15, 7)

######################################################################
# Constructing the circuit
# ------------------------
# We can now construct the corresponding single-mode Strawberry Fields program:
#

# Single-mode Strawberry Fields program
prog = sf.Program(1)

# Create the 7 Strawberry Fields free parameters for each layer
sf_params = []
names = ["r1", "sq_r", "sq_phi", "r2", "d_r", "d_phi", "kappa"]

for i in range(depth):
    # For the ith layer, generate parameter names "r1_i", "sq_r_i", etc.
    sf_params_names = ["{}_{}".format(n, i) for n in names]
    # Create the parameters, and append them to our list ``sf_params``.
    sf_params.append(prog.params(*sf_params_names))

######################################################################
# ``sf_params`` is now a nested list of shape ``(depth, 7)``, matching
# the shape of ``weights``.

sf_params = np.array(sf_params)
print(sf_params.shape)

######################################################################
# .. rst-class:: sphx-glr-script-out
#
#  Out:
#
#  .. code-block:: none
#
#     (15, 7)

######################################################################
# Now, we can create a function to define the :math:`i`\ th layer, acting
# on qumode ``q``. We add the :class:`~.utils.operation` decorator so that the layer can be used
# as a single operation when constructing our circuit within the usual
# Strawberry Fields Program context.

# layer architecture
@operation(1)
def layer(i, q):
    Rgate(sf_params[i][0]) | q
    Sgate(sf_params[i][1], sf_params[i][2]) | q
    Rgate(sf_params[i][3]) | q
    Dgate(sf_params[i][4], sf_params[i][5]) | q
    Kgate(sf_params[i][6]) | q
    return q

######################################################################
# We must also specify the input states to the variational quantum circuit
# - these are the Fock state :math:`\ket{i}`, :math:`i=0,\dots,d`,
# allowing us to optimize the circuit parameters to learn the target
# unitary acting on all input Fock states within the :math:`d`-dimensional
# subspace.
#

in_ket = np.zeros([gate_cutoff, cutoff])
np.fill_diagonal(in_ket, 1)

# Apply circuit of layers with corresponding depth
with prog.context as q:
    Ket(in_ket) | q
    for k in range(depth):
        layer(k) | q[0]

######################################################################
# Performing the optimization
# ---------------------------
#
# :math:`\newcommand{ket}[1]{\left|#1\right\rangle}`
# With the Strawberry Fields TensorFlow backend calculating the
# resulting state of the circuit symbolically, we can use TensorFlow to
# optimize the gate parameters to minimize the cost function we specify.
# With gate synthesis, we minimize the overlaps in the Fock basis
# between the target and learnt unitaries via the following cost function:
#
# .. math::
#
#     C(\vec{\theta}) = \frac{1}{d}\sum_{i=0}^{d-1} \left| \langle i \mid
#     V^\dagger U(\vec{\theta})\mid i\rangle - 1\right|
#
# where :math:`V` is the target unitary, :math:`U(\vec{\theta})` is the
# learnt unitary, and :math:`d` is the gate cutoff. Note that this is a
# generalization of state preparation to more than one input-output
# relation.
#
# For our target unitary, let's use Strawberry Fields to generate a 4x4
# random unitary:
#

from strawberryfields.utils import random_interferometer

# define unitary up to gate_cutoff
random_unitary = random_interferometer(gate_cutoff)
print(random_unitary)

# extend unitary up to cutoff
target_unitary = np.identity(cutoff, dtype=np.complex128)
target_unitary[:gate_cutoff, :gate_cutoff] = random_unitary

######################################################################
# .. rst-class:: sphx-glr-script-out
#
#  Out:
#
#  .. code-block:: none
#
#     [[ 0.23648826-0.48221431j  0.06829648+0.04447898j  0.51150074-0.09529866j
#        0.55205719-0.35974699j]
#      [-0.11148167+0.69780321j -0.24943828+0.08410701j  0.46705929-0.43192981j
#        0.16220654-0.01817602j]
#      [-0.22351926-0.25918352j  0.24364996-0.05375623j -0.09259829-0.53810588j
#        0.27267708+0.66941977j]
#      [ 0.11519953-0.28596729j -0.90164923-0.22099186j -0.09627758-0.13105595j
#       -0.0200152 +0.12766128j]]

######################################################################
# This matches the gate cutoff of :math:`d=4` that we chose above when
# defining our hyperparameters.
#
# Now, we instantiate the Strawberry Fields TensorFlow backend:

eng = sf.Engine('tf', backend_options={"cutoff_dim": cutoff, "batch_size": gate_cutoff})

######################################################################
# Here, we use the ``batch_size`` argument to perform the optimization in
# parallel - each batch calculates the variational quantum circuit acting
# on a different input Fock state: :math:`U(\vec{\theta}) | n\rangle`.
#

in_state = np.arange(gate_cutoff)

# extract action of the target unitary acting on
# the allowed input fock states.
target_kets = np.array([target_unitary[:, i] for i in in_state])
target_kets = tf.constant(target_kets, dtype=tf.complex64)

######################################################################
# Using this target unitary, we define the cost function we would like to
# minimize. We must use TensorFlow functions to manipulate this data, as well
# as a ``GradientTape`` to keep track of the corresponding gradients!
#
def cost(weights):
    # Create a dictionary mapping from the names of the Strawberry Fields
    # free parameters to the TensorFlow weight values.
    mapping = {p.name: w for p, w in zip(sf_params.flatten(), tf.reshape(weights, [-1]))}

    # Run engine
    state = eng.run(prog, args=mapping).state

    # Extract the statevector
    ket = state.ket()

    # overlaps
    overlaps = tf.math.real(tf.einsum('bi,bi->b', tf.math.conj(target_kets), ket))
    mean_overlap = tf.reduce_mean(overlaps)

    # Objective function to minimize
    cost = tf.abs(tf.reduce_sum(overlaps - 1))
    return cost, overlaps, ket

######################################################################
# Now that the cost function is defined, we can define and run the
# optimization. Below, we choose the Adam optimizer that is built into
# TensorFlow.
#

# Using Adam algorithm for optimization
opt = tf.keras.optimizers.Adam(learning_rate=lr)

######################################################################
# We then loop over all repetitions, storing the best predicted overlap
# value.
#

overlap_progress = []
cost_progress = []

# Run optimization
for i in range(reps):

    # reset the engine if it has already been executed
    if eng.run_progs:
        eng.reset()

    # one repetition of the optimization
    with tf.GradientTape() as tape:
        loss, overlaps_val, ket_val = cost(weights)

    # calculate the mean overlap
    # This gives us an idea of how the optimization is progressing
    mean_overlap_val = np.mean(overlaps_val)

    # store cost at each step
    cost_progress.append(loss)
    overlap_progress.append(overlaps_val)

    # one repetition of the optimization
    gradients = tape.gradient(loss, weights)
    opt.apply_gradients(zip([gradients], [weights]))

    # Prints progress at every rep
    if i % 1 == 0:
        # print progress
        print("Rep: {} Cost: {:.4f} Mean overlap: {:.4f}".format(i, loss, mean_overlap_val))

######################################################################
# .. rst-class:: sphx-glr-script-out
#
#  Out:
#
#  .. code-block:: none
#
#     Rep: 0 Cost: 4.2947 Mean overlap: -0.0737
#     Rep: 1 Cost: 3.2125 Mean overlap: 0.1969
#     Rep: 2 Cost: 4.3774 Mean overlap: -0.0944
#     Rep: 3 Cost: 3.6844 Mean overlap: 0.0789
#     Rep: 4 Cost: 3.7196 Mean overlap: 0.0701
#     Rep: 5 Cost: 3.2360 Mean overlap: 0.1910
#     Rep: 6 Cost: 3.1782 Mean overlap: 0.2055
#     Rep: 7 Cost: 3.1090 Mean overlap: 0.2228
#     Rep: 8 Cost: 3.0612 Mean overlap: 0.2347
#     Rep: 9 Cost: 3.0602 Mean overlap: 0.2349
#     Rep: 10 Cost: 2.7062 Mean overlap: 0.3234
#     Rep: 11 Cost: 2.7678 Mean overlap: 0.3081
#     Rep: 12 Cost: 2.2194 Mean overlap: 0.4451
#     Rep: 13 Cost: 2.0853 Mean overlap: 0.4787
#     Rep: 14 Cost: 2.2453 Mean overlap: 0.4387
#     Rep: 15 Cost: 1.6812 Mean overlap: 0.5797
#     Rep: 16 Cost: 1.9051 Mean overlap: 0.5237
#     Rep: 17 Cost: 1.4323 Mean overlap: 0.6419
#     Rep: 18 Cost: 1.4057 Mean overlap: 0.6486
#     Rep: 19 Cost: 1.2089 Mean overlap: 0.6978
#     Rep: 20 Cost: 1.1891 Mean overlap: 0.7027
#     Rep: 21 Cost: 1.2000 Mean overlap: 0.7000
#     Rep: 22 Cost: 1.1956 Mean overlap: 0.7011
#     Rep: 23 Cost: 1.1900 Mean overlap: 0.7025
#     Rep: 24 Cost: 1.1602 Mean overlap: 0.7100
#     Rep: 25 Cost: 1.0940 Mean overlap: 0.7265
#     Rep: 26 Cost: 1.0667 Mean overlap: 0.7333
#     Rep: 27 Cost: 1.0020 Mean overlap: 0.7495
#     Rep: 28 Cost: 0.9791 Mean overlap: 0.7552
#     Rep: 29 Cost: 0.9648 Mean overlap: 0.7588
#     Rep: 30 Cost: 0.9356 Mean overlap: 0.7661
#     Rep: 31 Cost: 0.9526 Mean overlap: 0.7618
#     Rep: 32 Cost: 0.8805 Mean overlap: 0.7799
#     Rep: 33 Cost: 0.9031 Mean overlap: 0.7742
#     Rep: 34 Cost: 0.8137 Mean overlap: 0.7966
#     Rep: 35 Cost: 0.8456 Mean overlap: 0.7886
#     Rep: 36 Cost: 0.8192 Mean overlap: 0.7952
#     Rep: 37 Cost: 0.8201 Mean overlap: 0.7950
#     Rep: 38 Cost: 0.8197 Mean overlap: 0.7951
#     Rep: 39 Cost: 0.7660 Mean overlap: 0.8085
#     Rep: 40 Cost: 0.7770 Mean overlap: 0.8058
#     Rep: 41 Cost: 0.7360 Mean overlap: 0.8160
#     Rep: 42 Cost: 0.7346 Mean overlap: 0.8163
#     Rep: 43 Cost: 0.7328 Mean overlap: 0.8168
#     Rep: 44 Cost: 0.7027 Mean overlap: 0.8243
#     Rep: 45 Cost: 0.7108 Mean overlap: 0.8223
#     Rep: 46 Cost: 0.6900 Mean overlap: 0.8275
#     Rep: 47 Cost: 0.6831 Mean overlap: 0.8292
#     Rep: 48 Cost: 0.6847 Mean overlap: 0.8288
#     Rep: 49 Cost: 0.6627 Mean overlap: 0.8343
#     Rep: 50 Cost: 0.6624 Mean overlap: 0.8344
#     Rep: 51 Cost: 0.6518 Mean overlap: 0.8370
#     Rep: 52 Cost: 0.6354 Mean overlap: 0.8412
#     Rep: 53 Cost: 0.6388 Mean overlap: 0.8403
#     Rep: 54 Cost: 0.6310 Mean overlap: 0.8422
#     Rep: 55 Cost: 0.6186 Mean overlap: 0.8453
#     Rep: 56 Cost: 0.6168 Mean overlap: 0.8458
#     Rep: 57 Cost: 0.6052 Mean overlap: 0.8487
#     Rep: 58 Cost: 0.5878 Mean overlap: 0.8531
#     Rep: 59 Cost: 0.5823 Mean overlap: 0.8544
#     Rep: 60 Cost: 0.5790 Mean overlap: 0.8552
#     Rep: 61 Cost: 0.5666 Mean overlap: 0.8583
#     Rep: 62 Cost: 0.5546 Mean overlap: 0.8614
#     Rep: 63 Cost: 0.5487 Mean overlap: 0.8628
#     Rep: 64 Cost: 0.5416 Mean overlap: 0.8646
#     Rep: 65 Cost: 0.5304 Mean overlap: 0.8674
#     Rep: 66 Cost: 0.5214 Mean overlap: 0.8696
#     Rep: 67 Cost: 0.5165 Mean overlap: 0.8709
#     Rep: 68 Cost: 0.5098 Mean overlap: 0.8726
#     Rep: 69 Cost: 0.4999 Mean overlap: 0.8750
#     Rep: 70 Cost: 0.4911 Mean overlap: 0.8772
#     Rep: 71 Cost: 0.4850 Mean overlap: 0.8788
#     Rep: 72 Cost: 0.4789 Mean overlap: 0.8803
#     Rep: 73 Cost: 0.4711 Mean overlap: 0.8822
#     Rep: 74 Cost: 0.4627 Mean overlap: 0.8843
#     Rep: 75 Cost: 0.4551 Mean overlap: 0.8862
#     Rep: 76 Cost: 0.4484 Mean overlap: 0.8879
#     Rep: 77 Cost: 0.4420 Mean overlap: 0.8895
#     Rep: 78 Cost: 0.4360 Mean overlap: 0.8910
#     Rep: 79 Cost: 0.4307 Mean overlap: 0.8923
#     Rep: 80 Cost: 0.4261 Mean overlap: 0.8935
#     Rep: 81 Cost: 0.4217 Mean overlap: 0.8946
#     Rep: 82 Cost: 0.4175 Mean overlap: 0.8956
#     Rep: 83 Cost: 0.4135 Mean overlap: 0.8966
#     Rep: 84 Cost: 0.4106 Mean overlap: 0.8974
#     Rep: 85 Cost: 0.4089 Mean overlap: 0.8978
#     Rep: 86 Cost: 0.4093 Mean overlap: 0.8977
#     Rep: 87 Cost: 0.4116 Mean overlap: 0.8971
#     Rep: 88 Cost: 0.4185 Mean overlap: 0.8954
#     Rep: 89 Cost: 0.4284 Mean overlap: 0.8929
#     Rep: 90 Cost: 0.4437 Mean overlap: 0.8891
#     Rep: 91 Cost: 0.4490 Mean overlap: 0.8877
#     Rep: 92 Cost: 0.4441 Mean overlap: 0.8890
#     Rep: 93 Cost: 0.4154 Mean overlap: 0.8962
#     Rep: 94 Cost: 0.3849 Mean overlap: 0.9038
#     Rep: 95 Cost: 0.3642 Mean overlap: 0.9089
#     Rep: 96 Cost: 0.3601 Mean overlap: 0.9100
#     Rep: 97 Cost: 0.3686 Mean overlap: 0.9079
#     Rep: 98 Cost: 0.3814 Mean overlap: 0.9047
#     Rep: 99 Cost: 0.3918 Mean overlap: 0.9020
#     Rep: 100 Cost: 0.3890 Mean overlap: 0.9027
#     Rep: 101 Cost: 0.3765 Mean overlap: 0.9059
#     Rep: 102 Cost: 0.3562 Mean overlap: 0.9110
#     Rep: 103 Cost: 0.3395 Mean overlap: 0.9151
#     Rep: 104 Cost: 0.3304 Mean overlap: 0.9174
#     Rep: 105 Cost: 0.3291 Mean overlap: 0.9177
#     Rep: 106 Cost: 0.3333 Mean overlap: 0.9167
#     Rep: 107 Cost: 0.3396 Mean overlap: 0.9151
#     Rep: 108 Cost: 0.3465 Mean overlap: 0.9134
#     Rep: 109 Cost: 0.3496 Mean overlap: 0.9126
#     Rep: 110 Cost: 0.3499 Mean overlap: 0.9125
#     Rep: 111 Cost: 0.3426 Mean overlap: 0.9144
#     Rep: 112 Cost: 0.3324 Mean overlap: 0.9169
#     Rep: 113 Cost: 0.3190 Mean overlap: 0.9203
#     Rep: 114 Cost: 0.3071 Mean overlap: 0.9232
#     Rep: 115 Cost: 0.2975 Mean overlap: 0.9256
#     Rep: 116 Cost: 0.2909 Mean overlap: 0.9273
#     Rep: 117 Cost: 0.2868 Mean overlap: 0.9283
#     Rep: 118 Cost: 0.2845 Mean overlap: 0.9289
#     Rep: 119 Cost: 0.2838 Mean overlap: 0.9290
#     Rep: 120 Cost: 0.2848 Mean overlap: 0.9288
#     Rep: 121 Cost: 0.2888 Mean overlap: 0.9278
#     Rep: 122 Cost: 0.2966 Mean overlap: 0.9258
#     Rep: 123 Cost: 0.3116 Mean overlap: 0.9221
#     Rep: 124 Cost: 0.3308 Mean overlap: 0.9173
#     Rep: 125 Cost: 0.3530 Mean overlap: 0.9118
#     Rep: 126 Cost: 0.3528 Mean overlap: 0.9118
#     Rep: 127 Cost: 0.3306 Mean overlap: 0.9173
#     Rep: 128 Cost: 0.2866 Mean overlap: 0.9283
#     Rep: 129 Cost: 0.2533 Mean overlap: 0.9367
#     Rep: 130 Cost: 0.2439 Mean overlap: 0.9390
#     Rep: 131 Cost: 0.2555 Mean overlap: 0.9361
#     Rep: 132 Cost: 0.2737 Mean overlap: 0.9316
#     Rep: 133 Cost: 0.2785 Mean overlap: 0.9304
#     Rep: 134 Cost: 0.2655 Mean overlap: 0.9336
#     Rep: 135 Cost: 0.2411 Mean overlap: 0.9397
#     Rep: 136 Cost: 0.2240 Mean overlap: 0.9440
#     Rep: 137 Cost: 0.2216 Mean overlap: 0.9446
#     Rep: 138 Cost: 0.2291 Mean overlap: 0.9427
#     Rep: 139 Cost: 0.2360 Mean overlap: 0.9410
#     Rep: 140 Cost: 0.2332 Mean overlap: 0.9417
#     Rep: 141 Cost: 0.2215 Mean overlap: 0.9446
#     Rep: 142 Cost: 0.2068 Mean overlap: 0.9483
#     Rep: 143 Cost: 0.1970 Mean overlap: 0.9507
#     Rep: 144 Cost: 0.1943 Mean overlap: 0.9514
#     Rep: 145 Cost: 0.1964 Mean overlap: 0.9509
#     Rep: 146 Cost: 0.1992 Mean overlap: 0.9502
#     Rep: 147 Cost: 0.1992 Mean overlap: 0.9502
#     Rep: 148 Cost: 0.1956 Mean overlap: 0.9511
#     Rep: 149 Cost: 0.1886 Mean overlap: 0.9529
#     Rep: 150 Cost: 0.1805 Mean overlap: 0.9549
#     Rep: 151 Cost: 0.1729 Mean overlap: 0.9568
#     Rep: 152 Cost: 0.1670 Mean overlap: 0.9582
#     Rep: 153 Cost: 0.1629 Mean overlap: 0.9593
#     Rep: 154 Cost: 0.1602 Mean overlap: 0.9599
#     Rep: 155 Cost: 0.1586 Mean overlap: 0.9603
#     Rep: 156 Cost: 0.1581 Mean overlap: 0.9605
#     Rep: 157 Cost: 0.1588 Mean overlap: 0.9603
#     Rep: 158 Cost: 0.1613 Mean overlap: 0.9597
#     Rep: 159 Cost: 0.1671 Mean overlap: 0.9582
#     Rep: 160 Cost: 0.1782 Mean overlap: 0.9554
#     Rep: 161 Cost: 0.1975 Mean overlap: 0.9506
#     Rep: 162 Cost: 0.2244 Mean overlap: 0.9439
#     Rep: 163 Cost: 0.2522 Mean overlap: 0.9369
#     Rep: 164 Cost: 0.2572 Mean overlap: 0.9357
#     Rep: 165 Cost: 0.2291 Mean overlap: 0.9427
#     Rep: 166 Cost: 0.1776 Mean overlap: 0.9556
#     Rep: 167 Cost: 0.1401 Mean overlap: 0.9650
#     Rep: 168 Cost: 0.1342 Mean overlap: 0.9664
#     Rep: 169 Cost: 0.1540 Mean overlap: 0.9615
#     Rep: 170 Cost: 0.1779 Mean overlap: 0.9555
#     Rep: 171 Cost: 0.1819 Mean overlap: 0.9545
#     Rep: 172 Cost: 0.1622 Mean overlap: 0.9595
#     Rep: 173 Cost: 0.1354 Mean overlap: 0.9662
#     Rep: 174 Cost: 0.1232 Mean overlap: 0.9692
#     Rep: 175 Cost: 0.1301 Mean overlap: 0.9675
#     Rep: 176 Cost: 0.1448 Mean overlap: 0.9638
#     Rep: 177 Cost: 0.1523 Mean overlap: 0.9619
#     Rep: 178 Cost: 0.1451 Mean overlap: 0.9637
#     Rep: 179 Cost: 0.1301 Mean overlap: 0.9675
#     Rep: 180 Cost: 0.1182 Mean overlap: 0.9704
#     Rep: 181 Cost: 0.1162 Mean overlap: 0.9709
#     Rep: 182 Cost: 0.1223 Mean overlap: 0.9694
#     Rep: 183 Cost: 0.1296 Mean overlap: 0.9676
#     Rep: 184 Cost: 0.1324 Mean overlap: 0.9669
#     Rep: 185 Cost: 0.1285 Mean overlap: 0.9679
#     Rep: 186 Cost: 0.1208 Mean overlap: 0.9698
#     Rep: 187 Cost: 0.1133 Mean overlap: 0.9717
#     Rep: 188 Cost: 0.1091 Mean overlap: 0.9727
#     Rep: 189 Cost: 0.1089 Mean overlap: 0.9728
#     Rep: 190 Cost: 0.1114 Mean overlap: 0.9721
#     Rep: 191 Cost: 0.1149 Mean overlap: 0.9713
#     Rep: 192 Cost: 0.1178 Mean overlap: 0.9706
#     Rep: 193 Cost: 0.1194 Mean overlap: 0.9702
#     Rep: 194 Cost: 0.1191 Mean overlap: 0.9702
#     Rep: 195 Cost: 0.1176 Mean overlap: 0.9706
#     Rep: 196 Cost: 0.1148 Mean overlap: 0.9713
#     Rep: 197 Cost: 0.1119 Mean overlap: 0.9720
#     Rep: 198 Cost: 0.1090 Mean overlap: 0.9728
#     Rep: 199 Cost: 0.1064 Mean overlap: 0.9734

######################################################################
# Results and visualisation
# -------------------------
#
# Plotting the cost vs. optimization step:
#

from matplotlib import pyplot as plt
# %matplotlib inline
plt.rcParams['font.family'] = 'serif'
plt.rcParams['font.sans-serif'] = ['Computer Modern Roman']
plt.style.use('default')

plt.plot(cost_progress)
plt.ylabel('Cost')
plt.xlabel('Step')
plt.show()

######################################################################
# .. image:: /_static/images/sphx_glr_run_gate_synthesis_001.png
#     :class: sphx-glr-single-img

######################################################################
# We can use matrix plots to plot the real and imaginary components of the
# target unitary :math:`V` and learnt unitary :math:`U`.
#

learnt_unitary = ket_val.numpy().T[:gate_cutoff, :gate_cutoff]
target_unitary = target_unitary[:gate_cutoff, :gate_cutoff]

fig, ax = plt.subplots(1, 4, figsize=(7, 4))
ax[0].matshow(target_unitary.real, cmap=plt.get_cmap('Reds'))
ax[1].matshow(target_unitary.imag, cmap=plt.get_cmap('Greens'))
ax[2].matshow(learnt_unitary.real, cmap=plt.get_cmap('Reds'))
ax[3].matshow(learnt_unitary.imag, cmap=plt.get_cmap('Greens'))

ax[0].set_xlabel(r'$\mathrm{Re}(V)$')
ax[1].set_xlabel(r'$\mathrm{Im}(V)$')
ax[2].set_xlabel(r'$\mathrm{Re}(U)$')
ax[3].set_xlabel(r'$\mathrm{Im}(U)$')
fig.show()

######################################################################
# .. image:: /_static/images/sphx_glr_run_gate_synthesis_002.png
#     :class: sphx-glr-single-img

######################################################################
# Process fidelity
# ----------------
#
# The process fidelity between the two unitaries is defined by
#
# .. math:: F_e  = \left| \left\langle \Psi(V) \mid \Psi(U)\right\rangle\right|^2
#
# where:
#
# -  :math:`\left|\Psi(V)\right\rangle` is the action of :math:`V` on one
#    half of a maximally entangled state :math:`\left|\phi\right\rangle`:
#
# .. math:: \left|\Psi(V)\right\rangle = (I\otimes V)\left|\phi\right\rangle,
#
# -  :math:`V` is the target unitary,
# -  :math:`U` the learnt unitary.
#

I = np.identity(gate_cutoff)
phi = I.flatten()/np.sqrt(gate_cutoff)
psiV = np.kron(I, target_unitary) @ phi
psiU = np.kron(I, learnt_unitary) @ phi

######################################################################
# Therefore, after 200 repetitions, the learnt unitary synthesized via a
# variational quantum circuit has the following process fidelity to the target
# unitary:
#

print(np.abs(np.vdot(psiV, psiU))**2)

######################################################################
# .. rst-class:: sphx-glr-script-out
#
#  Out:
#
#  .. code-block:: none
#
#     0.9481729779485119

######################################################################
# Circuit parameters
# ------------------
#
# We can also query the optimal variational circuit parameters
# :math:`\vec{\theta}` that resulted in the learnt unitary. For example,
# to determine the maximum squeezing magnitude in the variational quantum
# circuit:
#

print(np.max(np.abs(weights[:, 0])))

######################################################################
# .. rst-class:: sphx-glr-script-out
#
#  Out:
#
#  .. code-block:: none
#
#     0.39022577

######################################################################
# Further results
# ---------------
#
# After downloading the tutorial, even more refined results can be obtained by
# increasing the number of repetitions (``reps``), changing the depth of the
# circuit or altering the gate cutoff!
#

######################################################################
# References
# ----------
#
# 1. Juan Miguel Arrazola, Thomas R. Bromley, Josh Izaac, Casey R. Myers,
#    Kamil Brádler, and Nathan Killoran. Machine learning method for state
#    preparation and gate synthesis on photonic quantum computers. `Quantum
#    Science and Technology, 4
#    024004 <https://iopscience.iop.org/article/10.1088/2058-9565/aaf59e>`__,
#    (2019).
#
# 2. Nathan Killoran, Thomas R. Bromley, Juan Miguel Arrazola, Maria Schuld,
#    Nicolas Quesada, and Seth Lloyd. Continuous-variable quantum neural networks.
#    `Physical Review Research, 1(3), 033063.
#    <https://journals.aps.org/prresearch/abstract/10.1103/PhysRevResearch.1.033063>`__,
#    (2019).


text:
import matplotlib.pyplot as plt

import numpy as np
import seaborn as sns
import pandas as pd
import os
from os.path import join as pjoin
from sklearn.metrics import roc_curve
from sklearn.metrics import roc_auc_score

import matplotlib

font = {"size": 30}
matplotlib.rc("font", **font)
matplotlib.rcParams["text.usetex"] = True

DATA_DIR = "../out/cai_good"

files = os.listdir(DATA_DIR)

p_list = [10, 100, 1000]

plt.figure(figsize=(7, 7))

######## CPLVM ##########

# Load CPLVM results
# cplvm_stats = np.load(
#     pjoin(DATA_DIR, "bfs_targeted.npy")
# )
cplvm_stats = np.load(
    pjoin("/Users/andrewjones/Desktop", "bfs_targeted.npy")
)

bfs_stimulated_set = cplvm_stats[:, 0]
bfs_unstimulated_set = np.ndarray.flatten(cplvm_stats[:, 1:])
# import ipdb; ipdb.set_trace()

tpr_shuffled, fpr_shuffled, thresholds_shuffled = roc_curve(
    y_true=np.concatenate(
        [np.zeros(len(bfs_unstimulated_set)), np.ones(len(bfs_stimulated_set))]
    ),
    y_score=np.concatenate([bfs_unstimulated_set, bfs_stimulated_set]),
)
plt.plot(tpr_shuffled, fpr_shuffled, label="CPLVM", linestyle=":", color="black")
# plt.show()

cplvm_df = pd.DataFrame({"TPR": tpr_shuffled, "FPR": fpr_shuffled})
cplvm_df['method'] = "CPLVM"

######## Li 2012 ##########
li2012_stats_experiment = np.load(
    "../out/li2012/test_stats_experiment.npy"
)
li2012_stats_null = np.load("../out/li2012/test_stats_shuffled.npy")

tpr_shuffled, fpr_shuffled, thresholds_shuffled = roc_curve(
    y_true=np.concatenate(
        [np.zeros(len(li2012_stats_null)), np.ones(len(li2012_stats_experiment))]
    ),
    y_score=np.concatenate([li2012_stats_null, li2012_stats_experiment]),
)
plt.plot(tpr_shuffled, fpr_shuffled, label="Li 2012", linestyle="--", color="black")
li2012_df = pd.DataFrame({"TPR": tpr_shuffled, "FPR": fpr_shuffled})
li2012_df['method'] = "Li 2012"

# plt.legend(prop={"size": 20})
# plt.xlabel("TPR")
# plt.ylabel("FPR")
# plt.plot([0, 1], [0, 1], "--", color="black")

# plot_df = pd.concat([cplvm_df, li2012_df], axis=0)
# g = sns.lineplot(data=plot_df, x="TPR", y="FPR", style="method", color="black", ci=95, err_style="band")
# g.legend_.set_title(None)
plt.legend(prop={"size": 20})
plt.xlabel("TPR")
plt.ylabel("FPR")

# import ipdb; ipdb.set_trace()

plt.tight_layout()
plt.savefig("../out/roc_comparison_genesets.png")
plt.show()
import ipdb; ipdb.set_trace()

max_stars_repo_name: andrewcharlesjones/cplvm

max_issues_repo_name: andrewcharlesjones/cplvm

max_forks_repo_name: andrewcharlesjones/cplvm


text:
import abc
import logging

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn as nn

import rlkit.torch.pytorch_util as ptu
from rlkit.policies.base import ExplorationPolicy
from rlkit.torch.core import torch_ify, elem_or_tuple_to_numpy
from rlkit.torch.distributions import (
    Delta, TanhNormal, MultivariateDiagonalNormal, GaussianMixture, GaussianMixtureFull,
)
from rlkit.torch.networks import Mlp, CNN, GAT
from rlkit.torch.networks.basic import MultiInputSequential
from rlkit.torch.networks.stochastic.distribution_generator import (
    DistributionGenerator
)

class TorchStochasticPolicy(
    DistributionGenerator,
    ExplorationPolicy, metaclass=abc.ABCMeta
):
    def get_action(self, obs_np, ):
        actions = self.get_actions(obs_np[None])
        return actions[0, :], {}

    def get_actions(self, obs_np, ):
        dist = self._get_dist_from_np(obs_np)
        actions = dist.sample()
        return elem_or_tuple_to_numpy(actions)

    def _get_dist_from_np(self, *args, **kwargs):
        torch_args = tuple(torch_ify(x) for x in args)
        torch_kwargs = {k: torch_ify(v) for k, v in kwargs.items()}
        dist = self(*torch_args, **torch_kwargs)
        return dist

class TorchGATStochasticPolicy(
    DistributionGenerator,
    ExplorationPolicy, metaclass=abc.ABCMeta
):
    def get_action(self, obs_np, ):
        actions = self.get_actions(obs_np)
        return actions[0, :], {}

    def get_actions(self, obs_np, ):
        dist = self._get_dist_from_np(obs_np)
        actions = dist.sample()
        return elem_or_tuple_to_numpy(actions)

    def _get_dist_from_np(self, *args, **kwargs):
        torch_args = tuple(torch_ify(x) for x in args)
        torch_kwargs = {k: torch_ify(v) for k, v in kwargs.items()}
        dist = self(*torch_args, **torch_kwargs)
        return dist

class PolicyFromDistributionGenerator(
    MultiInputSequential,
    TorchStochasticPolicy,
):
    """
    Usage:
    ```
    distribution_generator = FancyGenerativeModel()
    policy = PolicyFromBatchDistributionModule(distribution_generator)
    ```
    """
    pass

class MakeDeterministic(TorchStochasticPolicy):
    def __init__(
            self,
            action_distribution_generator: DistributionGenerator,
    ):
        super().__init__()
        self._action_distribution_generator = action_distribution_generator

    def forward(self, *args, **kwargs):
        dist = self._action_distribution_generator.forward(*args, **kwargs)
        return Delta(dist.mle_estimate())

class MakeGATDeterministic(TorchGATStochasticPolicy):
    def __init__(
            self,
            action_distribution_generator: DistributionGenerator,
    ):
        super().__init__()
        self._action_distribution_generator = action_distribution_generator

    def forward(self, *args, **kwargs):
        dist = self._action_distribution_generator.forward(*args, **kwargs)
        return Delta(dist.mle_estimate())

max_stars_repo_name: guilhermecano/rlk

max_issues_repo_name: guilhermecano/rlk

max_forks_repo_name: guilhermecano/rlk


text:
from model import model
import numpy as np
import os
from mcmcpy.mcmc.mcmc_sampler import MCMCSampler

def test_model():
    """
    Initial test to ensure model has not been altered.
    """
    # set up ground truth
    a = 2
    b = 3.5
    x = [1, 2, 3, 4, 5, 6]
    y = np.array(x) * a + b

    # instance model / evaluate
    m = model(x)
    y_model = m.evaluate({'a': a, 'b': b})

    # test
    np.testing.assert_equal(y, y_model)

def test_mcmc_fit():
    """
    Tests the fit method of the mcmc class.
    """
    # set up ground truth
    a = 2
    b = 3.5
    x = [1, 2, 3, 4, 5, 6]
    y = np.array(x) * a + b

    # instance model
    m = model(x)

    # define priors for params
    params = {'a': ['Uniform', -50.0, 50.0],
              'b': ['Uniform', -50.0, 50.0]}

    # initial param guess
    params0 = {'a': 0.0, 'b': 0.0}

    # instance prognosis class
    mcmc = MCMCSampler(data=y, model=m, params=params)

    # LSQ fit
    params_opt, ssq_opt = mcmc.fit(params0, repeats=2, opt_method='L-BFGS-B')

    # test params
    params_true = [a, b]
    params_fit = [params_opt['a'], params_opt['b']]
    np.testing.assert_array_almost_equal(params_true, params_fit, 4)

    # test output
    y_opt = m.evaluate(params_opt)
    np.testing.assert_array_almost_equal(y, y_opt, 5)

def test_mcmc_standard():
    """
    Test parameter estimation using standard (non-adaptive) step method.
    """
    # set up ground truth / instance model
    a = 2
    b = 3.5
    x = np.arange(200)
    m = model(x)
    std_dev = 0.01
    yn = m.generate_noisy_data({'a': a, 'b': b}, std_dev)

    # define priors for params
    params = {'a': ['Uniform', -50.0, 50.0],
              'b': ['Uniform', -50.0, 50.0]}

    # instance prognosis class
    mcmc = MCMCSampler(data=yn, model=m, params=params)

    # LSQ fit for get good starting point
    params_opt, ssq_opt = ({'a': 2.00052, 'b': 3.499}, 5.494477e-05)

    # generate pymc model
    mcmc.generate_pymc_model(q0=params_opt, ssq0=ssq_opt)

    # sample
    N = 1E4
    burn = 2E3
    mcmc.sample(N, burn, step_method='metropolis',
                scales={'a': 0.015, 'b': 0.015})

    # extract chain traces
    a_trace = mcmc.MCMC.trace('a')[:]
    b_trace = mcmc.MCMC.trace('b')[:]
    std_dev_trace = mcmc.MCMC.trace('std_dev')[:]

    # test
    np.testing.assert_almost_equal(np.mean(a_trace), a, 2)
    np.testing.assert_almost_equal(np.mean(b_trace), b, 2)

    # remove unwanted files
    os.remove('mcmc.p')

def test_mcmc_adaptive():
    """
    Test parameter estimation using standard (non-adaptive) step method.
    """
    # set up ground truth / instance model
    a = 2
    b = 3.5
    x = np.arange(200)
    m = model(x)
    std_dev = 0.01
    yn = m.generate_noisy_data({'a': a, 'b': b}, std_dev)

    # define priors for params
    params = {'a': ['Uniform', -50.0, 50.0],
              'b': ['Uniform', -50.0, 50.0]}

    # initial param guess
    params0 = {'a': 0.0, 'b': 0.0}

    # instance prognosis class
    mcmc = MCMCSampler(data=yn, model=m, params=params)

    # LSQ fit for get good starting point
    params_opt, ssq_opt = ({'a': 2.00052, 'b': 3.499}, 5.494477e-05)

    # generate pymc model
    mcmc.generate_pymc_model(q0=params_opt, ssq0=ssq_opt)

    # sample
    N = 1E4
    burn = 2E3
    mcmc.sample(N, burn, step_method='adaptive', interval=300,
                delay=burn/2, scales={'a': 0.015, 'b': 0.015})

    # extract chain traces
    a_trace = mcmc.MCMC.trace('a')[:]
    b_trace = mcmc.MCMC.trace('b')[:]
    std_dev_trace = mcmc.MCMC.trace('std_dev')[:]

    # make sure std_dev is actually being sampled
    if len(set(std_dev_trace)) < 2:
        raise ValueError('std_dev trace contains only one unique value!')

    # test
    np.testing.assert_almost_equal(a, np.mean(a_trace), 2)
    np.testing.assert_almost_equal(b, np.mean(b_trace), 2)
    np.testing.assert_almost_equal(std_dev, np.mean(std_dev_trace), 2)

    # remove unwanted files
    os.remove('mcmc.p')

max_stars_repo_name: nasa/MCMCPy

max_issues_repo_name: nasa/MCMCPy

max_forks_repo_name: nasa/MCMCPy


text:
import pickle
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf
import numpy as np
import time
from datetime import datetime
from metrics import recall_2at1, recall_at_k, precision_at_k, MRR, MAP
from data_utils_record import get_record_parser, get_batch_dataset
from nvidia_helper import get_available_gpu
from config import config
from model_FLS import model

data_path='ubuntu'
if __name__=="__main__":
    FLAGS = config(data_path)
    if FLAGS.auto_gpu:
        index_of_gpu = get_available_gpu()
        FLAGS.gpu = 'gpu:' + str(index_of_gpu)
        os.environ["CUDA_VISIBLE_DEVICES"] = FLAGS.gpu.split(':')[1]
    else:
        os.environ["CUDA_VISIBLE_DEVICES"] = '0'

    # Checkpoint directory.
    out_dir = os.path.abspath(os.path.join(os.path.curdir, FLAGS.log_root))
    checkpoint_dir = os.path.abspath(os.path.join(out_dir, "checkpoints/"))

    with tf.device("/%s" % FLAGS.gpu):
        session_conf = tf.ConfigProto(
            allow_soft_placement=FLAGS.allow_soft_placement,
            log_device_placement=FLAGS.log_device_placement)
        sess = tf.Session(config=session_conf)

        with sess.as_default():
            parser = get_record_parser(FLAGS)
            test_dataset = get_batch_dataset(FLAGS.test_record_file, parser, FLAGS.batch_size, FLAGS.num_threads, FLAGS.capacity, True)
            test_iterator = test_dataset.make_initializable_iterator()
            sess.run(test_iterator.initializer)

            test_handle = sess.run(test_iterator.string_handle())
            handle = tf.placeholder(tf.string, shape=[])
            iterator = tf.data.Iterator.from_string_handle(handle, test_dataset.output_types, test_dataset.output_shapes)

            model = model(iterator, FLAGS, FLAGS.embed_dim, FLAGS.vocab_size, FLAGS.char_embed_dim, FLAGS.char_vocab_size, FLAGS.rnn_dim, FLAGS.max_turn,
                            FLAGS.max_utterance_len, FLAGS.max_word_len)

            global_step = tf.Variable(0, name="global_step", trainable=False)

            # Restore all variables
            saver = tf.train.Saver()
            print(tf.train.latest_checkpoint(checkpoint_dir))
            saver.restore(sess, tf.train.latest_checkpoint(checkpoint_dir))

            def dev_step():
                acc = []
                losses = []
                pred_scores = []
                ture_scores = []
                count = 0
                while True:
                    try:
                        feed_dict = {
                            handle: test_handle,
                            model.dropout_keep_prob: 1.0
                        }
                        step, loss, accuracy, y_pred, target = sess.run(
                            [global_step, model.loss, model.accuracy, model.y_pred, model.target], feed_dict)

                        acc.append(accuracy)
                        losses.append(loss)
                        pred_scores += list(y_pred[:, 1])
                        ture_scores += list(target)

                        count +=1
                        if count % 1000 == 0:
                            print("Processing {} samples".format(count))
                    except tf.errors.OutOfRangeError:
                        break

                MeanAcc = sum(acc) / len(acc)
                MeanLoss = sum(losses) / len(losses)

                with open(os.path.join(out_dir, 'predScores-iter-%s-test.txt'%(step)), 'w') as f:
                    for score1, score2 in zip(pred_scores, ture_scores):
                        f.writelines(str(score1) + '\t' + str(score2) + '\n')

                num_sample = int(len(pred_scores) / 10)
                score_list = np.split(np.array(pred_scores), num_sample, axis=0)
                recall_2_1 = recall_2at1(score_list, k=1)

                recall_at_1 = recall_at_k(np.array(ture_scores),  np.array(pred_scores), 1)
                recall_at_2 = recall_at_k(np.array(ture_scores),  np.array(pred_scores), 2)
                recall_at_5 = recall_at_k(np.array(ture_scores),  np.array(pred_scores), 5)
                time_str = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                print("**********************************")
                print('pred_scores: ', len(pred_scores))
                print("recall_2_1:  %.3f" % (recall_2_1))
                print("recall_at_1: %.3f" % (recall_at_1))
                print("recall_at_2: %.3f" % (recall_at_2))
                print("recall_at_5: %.3f" % (recall_at_5))
                print("**********************************")

            dev_step()

max_stars_repo_name: chongyangtao/MRFN

max_issues_repo_name: chongyangtao/MRFN

max_forks_repo_name: chongyangtao/MRFN


text:
#INPUT
Tmin=-100.
Tmax=100.
THmin =-68. #because we don't start at standard pressure -86 is ~70C moist adiabat
THmax= 42.
Pbot = 105
Ptop = 1    #kPa - upper atmosphere limit surface
Plim = 1
degree =10      #degree of polinomial to model the curves
#=======================================================
#supress warnings
import warnings
def fxn():
    warnings.warn("deprecated", DeprecationWarning)
with warnings.catch_warnings():
    warnings.simplefilter("ignore")
    fxn()

#import necessary packages
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit, minimize
from scipy import spatial
import os.path

P0 = 100.      #kPa

#create temp and pressure axes
Pbrange = np.arange(Pbot,10,-0.0001)
Pmrange = np.arange(10,2,-0.00001)
Ptrange = np.arange(2,Ptop,-0.000001)
PrangeList = np.concatenate((Pbrange,Pmrange,Ptrange))

# PrangeList = np.concatenate((Pbrange,Pmrange))
Prange = PrangeList[:-1]

#create temp and pressure axes
nTw = np.arange(Tmin,Tmax,0.5)
nThetaW = np.arange(THmin,THmax,0.5)

datafileB = '%s-%s_B.npy' %(Tmin,Tmax)
datafileT = '%s-%s_T.npy' %(Tmin,Tmax)

dataB = np.load(datafileO)
dataT = np.load(datafileT)

#monotonically select points every 0.1 kPa for fitting
PrangeIdx = [np.argmin(abs(PrangeList - i)) for i in np.arange(Pbot,Plim,-0.1)]
PrangeFit = Prange[PrangeIdx]
P0idx = np.argmin(abs(Prange - P0))
Tidx = [np.argmin(abs(nTw - i)) for i in nThetaW]

subarrayTHwB = dataB[Tidx,:]
subarrayTHwT = dataT[Tidx,:]

fitarrayB = subarrayTHwB[:,PrangeIdx]
fitarrayT = subarrayTHwT[:,PrangeIdx]

diff = fitarrayT - fitarrayB
truthMAE = MAE = np.mean(abs(diff.ravel()))
print('TRUTH MAE: %s' %truthMAE)

plt.imshow(diff.T,aspect='auto',origin='lower',cmap='cubehelix_r',vmin=0,vmax=0.08)
plt.xlabel("$\\theta_w$ [$^\circ$C]")
plt.ylabel("pressure [kPa]")
ax = plt.gca()
ax.set_xticks(np.arange(0,len(nThetaW),20))
ax.set_xticklabels(C0axis[::20].astype(int))
ax.set_yticks(PaxisIdx)
ax.set_yticklabels(PrangeFit[PaxisIdx].round())
cbar = plt.colorbar(format='%.2f')
cbar.set_label('temperature difference [K]')
plt.title('SENSITIVITY TO CONSTANTS')
plt.savefig('./figs/TruthSensitivity.pdf')
plt.show()

reason: import numpy,from scipy


text:
"""
This example module shows some simple methods and best practices for documentation
"""

import numpy as np

def hello_world() -> str:
    """
    Method description - A simple method to get the hello world string

    Returns:
       The string "Hello World"
    """
    return "Hello World"

def add_value_to_numpy(array: np.ndarray, amount: float = 1) -> np.ndarray:
    """
    A sample method to add a value to every element in a pandas DataFrame.

    Args:
        array: The source DataFrame to work on.
        amount: The amount to add to each element in the DataFrame

    Returns:
        A new DataFrame with each value increased by amount.

    Examples:
        Examples should be written in doctest format, and should illustrate how
        to use the function.
        >>> array = np.array([1, 1, 1 ,1, 1])
        >>> result_array = add_value_to_numpy(array, 1)

    """
    if array is None or \
            not isinstance(array, np.ndarray):
        raise ValueError("array must be a valid ndarray")
    # if isinstance(a, np.ndarray):

    return array + amount

max_stars_repo_name: berndverst/data-science-template

max_issues_repo_name: berndverst/data-science-template

max_forks_repo_name: berndverst/data-science-template


text:
#!/usr/bin/env python
# encoding: utf-8
#
# Copyright SAS Institute
#
#  Licensed under the Apache License, Version 2.0 (the License);
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#

''' Transforms for ONNX graph '''

import numpy as np
import onnx
from onnx import helper, numpy_helper, shape_inference, mapping
from onnx import AttributeProto, TensorProto, GraphProto

from dlpy.model_conversion.onnx_graph import OnnxNode

class OpTypePattern(object):
    '''
    A tree pattern of operators to match in an ONNX graph

    Parameters
    ----------
    op_type : str
        Specifies the op type.
    name : str, optional
        Specifies a name for the node.
    outputs : list of :class:`OpTypePattern` objects, optional
        Specifies the output nodes.

    Returns
    -------
    :class:`OpTypePattern` object

    '''
    def __init__(self, op_type, name=None, outputs=None):
        self._op_type = op_type
        self._name = name
        if outputs is None:
            outputs = []
        self._outputs = [
            output_pattern if isinstance(output_pattern, OpTypePattern) else
            OpTypePattern(output_pattern) for output_pattern in outputs
        ]

    @property
    def op_type(self):
        return self._op_type

    @property
    def outputs(self):
        return self._outputs

    @property
    def name(self):
        return self._name

class Transformer(object):
    '''
    Transforms an OnnxGraph

    Parameters
    ----------
    pattern : :class:`OpTypePattern` object, optional
        The pattern to match.

    Returns
    -------
    :class:`Transformer` object

    '''
    def __init__(self, pattern=None):
        self.pattern = pattern

    def match(self, node, pattern=None):
        '''
        Checks if a subgraph rooted at `node` matches pattern.
        If there is a match, returns True.  If not, returns False.

        Parameters
        ----------
        node : :class:`OnnxNode` object
            The root node to be checked.
        pattern : :class:`OpTypePattern` object, optional
            The pattern to match.  If None, defaults to self.pattern.

        Returns
        -------
        boolean

        '''
        if pattern is None:
            if self.pattern is None:
                raise ValueError('No pattern to match.')
            pattern = self.pattern

        if node.op_type != pattern.op_type:
            return False
        elif pattern.outputs and len(node.children) != len(pattern.outputs):
            return False
        else:
            ret = []
            for child, child_pattern in zip(node.children, pattern.outputs):
                r = self.match(child, child_pattern)
                ret.append(r)
            return all(ret)

    def get_mapping(self, node, pattern=None):
        '''
        Given that `node` is the root of a matched subgraph, returns a dict
        mapping names of the OpTypePatterns to their matched OnnxNodes

        Parameters
        ----------
        node : :class:`OnnxNode` object
            The root node of a matching subgraph.
        pattern : :class:`OpTypePattern` object, optional
            The matching pattern.  If None, defaults to self.pattern.

        Returns
        -------
        dict
            key, value of OpTypePattern name and OnnxNode

        '''
        if pattern is None:
            if self.pattern is None:
                raise ValueError('No pattern to match.')
            pattern = self.pattern

        mapping_dict = {}
        def _mapping(node, pattern, mapping_dict):
            if pattern.name is None:
                raise ValueError('Cannot generate mapping dict,'
                                ' OpTypePattern name is None.')
            mapping_dict[pattern.name] = node
            for child, child_pattern in zip(node.children, pattern.outputs):
                _mapping(child, child_pattern, mapping_dict)
            return mapping_dict

        return _mapping(node, pattern, mapping_dict)

    def is_eligible(self, graph, node):
        '''
        Checks whether subgraph rooted at node is eligible for the transform.
        Each subclass should implement this.

        Parameters
        ----------
        graph : :class:`OnnxGraph` object
            The graph to be transformed.
        node : :class:`OnnxNode` object
            The root node of the subgraph to be transformed.

        Returns
        -------
        boolean

        '''
        return True

    def run_transform(self, graph, node):
        '''
        Define the transform for a single subgraph. Implemented by subclass.

        Parameters
        ----------
        graph : :class:`OnnxGraph` object
            The graph to be transformed.
        node : :class:`OnnxNode` object
            The root node of the subgraph to be transformed.

        Returns
        -------
        :class:`OnnxGraph` object
            The transformed graph

        '''
        return graph

    def __call__(self, graph):
        '''
        Call on `graph` to execute the transform on all eligible subgraphs

        Parameters
        ----------
        graph : :class:`OnnxGraph` object
            The graph to be transformed.

        Returns
        -------
        :class:`OnnxGraph` object
            The transformed graph

        '''
        matches = filter(lambda x: self.match(x), graph.node)
        ops = filter(lambda x: self.is_eligible(graph, x), matches)
        for op in ops:
            self.run_transform(graph, op)
            graph.connect_nodes()
        return graph

class ConstToInitializer(Transformer):
    ''' Remove constant ops and add tensor to initializer'''
    def __init__(self):
        pattern = OpTypePattern('Constant')
        super(ConstToInitializer, self).__init__(pattern)

    def run_transform(self, graph, node):
        tensor = numpy_helper.to_array(node.attrs['value'])
        graph.tensor_dict[node.output[0]] = tensor

        # remove the constant op
        graph.remove_node(node.name)

        return graph

class InitReshape(Transformer):
    ''' Remove reshape ops and add reshaped tensor to initializer'''
    def __init__(self):
        pattern = OpTypePattern('Reshape')
        super(InitReshape, self).__init__(pattern)

    def is_eligible(self, graph, node):
        if node.input[0] in graph.tensor_dict:
            return True
        return False

    def _get_shape(self, graph, node):
        ''' Get reshape op's shape '''
        name = node.input[1]
        shape = [int(x) for x in graph.tensor_dict[name].flatten()]
        return tuple(shape)

    def run_transform(self, graph, node):
        shape = self._get_shape(graph, node)
        tensor = graph.tensor_dict[node.input[0]]
        graph.tensor_dict[node.output[0]] = tensor.reshape(shape)

        # remove reshape op
        graph.remove_node(node.name)

        return graph

class InitUnsqueeze(Transformer):
    ''' Remove unsqueeze ops and add unsqueezed tensor to initializer'''
    def __init__(self):
        pattern = OpTypePattern('Unsqueeze')
        super(InitUnsqueeze, self).__init__(pattern)

    def is_eligible(self, graph, node):
        if node.input[0] in graph.tensor_dict:
            return True
        return False

    def _unsqueeze(self, tensor, axes):
        ''' unsqueeze tensor by specifying axes to be inserted '''
        _shape = list(tensor.shape)
        new_dim = len(tensor.shape) + len(axes)
        unsqueezed_shape = [1] * new_dim
        for i in range(new_dim):
            if i not in axes:
                unsqueezed_shape[i] = _shape.pop(0)

        return tensor.reshape(tuple(unsqueezed_shape))

    def run_transform(self, graph, node):
        axes = node.attrs['axes']
        tensor = graph.tensor_dict[node.input[0]]
        graph.tensor_dict[node.output[0]] = self._unsqueeze(tensor, axes)

        # remove unsqueeze op
        graph.remove_node(node.name)

        return graph

class FuseMulAddBN(Transformer):
    ''' Fuse Mul + Add into BN '''
    def __init__(self):
        add = OpTypePattern('Add', name='add')
        mul = OpTypePattern('Mul', name='mul', outputs=[add])
        bn = OpTypePattern('BatchNormalization', name='bn', outputs=[mul])
        super(FuseMulAddBN, self).__init__(bn)

    def is_eligible(self, graph, node):
        mapping = self.get_mapping(node)
        bn, mul, add = mapping['bn'], mapping['mul'], mapping['add']

        # only spatial batchnorm is supported
        if bn.attrs.get('spatial') is not None and bn.attrs['spatial'] != 1:
            return False
        # mul and add must be initialized by some tensor
        if (mul.input[0] not in graph.tensor_dict and
            mul.input[1] not in graph.tensor_dict):
            return False
        if (add.input[0] not in graph.tensor_dict and
            add.input[1] not in graph.tensor_dict):
            return False

        t = graph.tensor_dict
        scale = t[bn.input[1]]
        bias = t[bn.input[2]]
        _mul_tensor = t.get(mul.input[0], t[mul.input[1]])
        mul_tensor = np.squeeze(_mul_tensor)
        _add_tensor = t.get(add.input[0], t[add.input[1]])
        add_tensor = np.squeeze(_add_tensor)

        # check mul is broadcastable
        if mul_tensor.shape != scale.shape or mul_tensor.shape != bias.shape:
            if mul_tensor.shape != (1,) and mul_tensor.shape != ():
                return False

        # check add is broadcastable
        if add_tensor.shape != bias.shape:
            if add_tensor.shape != (1,) and add_tensor.shape != ():
                return False

        return True

    def run_transform(self, graph, node):
        mapping = self.get_mapping(node)
        bn, mul, add = mapping['bn'], mapping['mul'], mapping['add']

        t = graph.tensor_dict
        scale = t[bn.input[1]]
        bias = t[bn.input[2]]
        _mul_tensor = t.get(mul.input[0], t[mul.input[1]])
        mul_tensor = np.squeeze(_mul_tensor)
        _add_tensor = t.get(add.input[0], t[add.input[1]])
        add_tensor = np.squeeze(_add_tensor)

        # multiply scale and bias
        t[bn.input[1]] = np.multiply(scale, mul_tensor)
        _bias = np.multiply(bias, mul_tensor)
        t[bn.input[2]] = np.add(_bias, add_tensor)

        # connect output of bn to output of add
        bn.output[0] = add.output[0]

        # remove mul and add nodes
        graph.remove_node(mul.name)
        graph.remove_node(add.name)

        return graph