<commit_before>import numpy as np
import floq.parametric_system as ps
import floq.evolution as ev


class OptimizationTaskBase(object):
    """
    Base class to specify an optimization problem,
    i.e. a physical (parametrised) system,
    a fidelity, a target and an initial control vector.
    """
    def system(self):
        raise NotImplementedError

    def fidelity(self, control, t):
        raise NotImplementedError

    def grad_fidelity(self, control, t):
        raise NotImplementedError

    def target_fidelity(self, control, t):
        raise NotImplementedError

    def target(self):
        raise NotImplementedError

    def target_fid(self):
        raise NotImplementedError


class OptimizationTaskWithFunctions(OptimizationTaskBase):
    def __init__(self, system, fid, dfid, target, target_fid, init):
        """
        system: Instance of ParametricSystemBase
        fid: callable fid(fixed_system,u,target)
        dfid: callable fid(fixed_system,u,du,target)
        """

        self.system = system
        self.fid = fid
        self.dfid = dfid
        self.target = target
        self.target_fid = target_fid
        self.init = init

    def fidelity(self, controls, t):
        fixed = self.system.get_system(controls, t)
        u = ev.evolve_system(fixed)
        return self.fid(fixed, u, self.target)

    def grad_fidelity(self, controls, t):
        fixed = self.system.get_system(controls, t)
        u, du = ev.evolve_system_with_derivatives(fixed)
        return self.dfid(fixed, u, du, self.target)<commit_msg>Begin implementation of an OptimizationTask for ensembles<commit_after>import numpy as np
import floq.parametric_system as ps
import floq.evolution as ev


class OptimizationTaskBase(object):
    """
    Base class to specify an optimization problem,
    i.e. a physical (parametrised) system,
    a fidelity, a target and an initial control vector.
    """
    def fidelity(self, control, t):
        raise NotImplementedError

    def grad_fidelity(self, control, t):
        raise NotImplementedError

    def target_fidelity(self, control, t):
        raise NotImplementedError

    def target(self):
        raise NotImplementedError

    def target_fid(self):
        raise NotImplementedError


class OptimizationTaskWithFunctions(OptimizationTaskBase):
    def __init__(self, system, fid, dfid, target, target_fid, init):
        """
        system: Instance of ParametricSystemBase
        fid: callable fid(fixed_system,u,target)
        dfid: callable fid(fixed_system,u,du,target)
        """

        self.system = system
        self.fid = fid
        self.dfid = dfid
        self.target = target
        self.target_fid = target_fid
        self.init = init

    def fidelity(self, controls, t):
        fixed = self.system.get_system(controls, t)
        u = ev.evolve_system(fixed)
        return self.fid(fixed, u, self.target)

    def grad_fidelity(self, controls, t):
        fixed = self.system.get_system(controls, t)
        u, du = ev.evolve_system_with_derivatives(fixed)
        return self.dfid(fixed, u, du, self.target)


class EnsembleOptimizationTask(OptimizationTaskBase):
    """
    An optimization task where an ensemble of (non-interacting)
    systems is optimised.
    """

    def __init__(self, ensemble, fid, target, init):
        self.ensemble = ensemble
        self.fid = fid
        self.dfid = None
        self.target = target
        self.init = init

    def fidelity(self, controls, t):
        fixed_systems = self.ensemble.get_systems(controls, t)
        fid = 0.0
        for sys in fixed_systems:
            u = ev.evolve_system(sys)
            fid += self.fid(sys, u, self.target)
        return fid/self.ensemble.n
{'fraction_non_alphanumeric': 0.08558438085049479, 'fraction_numerical': 0.0005349023803155924, 'mean_word_length': 2.692003948667325, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6782474', 'n_tokens_mistral': 1146, 'n_tokens_neox': 1103, 'n_words': 301}