<commit_before>import numpy as np
from abc import ABC, abstractmethod
from inspect import getfullargspec

from desc import config
from desc.backend import use_jax, jnp, jit
from desc.utils import Timer
from desc.io import IOAble
from desc.derivatives import Derivative
from desc.compute import arg_order

# XXX: could use `indicies` instead of `arg_order` in ObjectiveFunction loops


class ObjectiveFunction(IOAble):
    """Objective function comprised of one or more Objectives."""

    _io_attrs_ = ["objectives", "constraints"]

    def __init__(self, objectives, constraints=(), eq=None, use_jit=True):
        """Initialize an Objective Function.

        Parameters
        ----------
        objectives : tuple of Objective
            List of objectives to be targeted during optimization.
        constraints : tuple of Objective, optional
            List of objectives to be used as constraints during optimization.
        eq : Equilibrium, optional
            Equilibrium that will be optimized to satisfy the Objective.
        use_jit : bool, optional
            Whether to just-in-time compile the objective and derivatives.

        """
        if not isinstance(objectives, tuple):
            objectives = (objectives,)
        if not isinstance(constraints, tuple):
            constraints = (constraints,)

        self._objectives = objectives
        self._constraints = constraints
        self._use_jit = use_jit
        self._built = False
        self._compiled = False

        if eq is not None:
            self.build(eq, use_jit=self._use_jit)

    def _set_state_vector(self):
        """Set state vector components, dimensions, and indicies."""
        self._args = np.concatenate([obj.args for obj in self.objectives])
        if self.constraints:
            self._args = np.concatenate(
                (self.args, np.concatenate([obj.args for obj in self.constraints]))
            )
        self._args = np.unique(self.args)

        self._dimensions = self.objectives[0].dimensions

        idx = 0
        self._indicies = {}
        for arg in arg_order:
            if arg in self.args:
                self._indicies[arg] = np.arange(idx, idx + self.dimensions[arg])
                idx += self.dimensions[arg]
            else:
                self._indicies[arg] = np.array([])

        self._dim_y = idx

    def _build_linear_constraints(self):
        """Compute and factorize A to get pseudoinverse and nullspace."""
        # A*y = b = B*y
        self._A = np.array([[]])
        self._B = np.array([[]])
        self._b = np.array([])
        for obj in self.constraints:
            A = np.array([[]])
            for arg in arg_order:
                if arg in self.args:
                    a = np.atleast_2d(obj.derivatives[arg])
                    A = np.hstack((A, a)) if A.size else a
            if obj.target_arg in self.args:
                B = np.eye(self.dim_y)[self.indicies[obj.target_arg], :]
            else:
                B = np.zeros((obj.dim_f, self.dim_y))
            b = obj.target
            self._A = np.vstack((self.A, A)) if self.A.size else A
            self._B = np.vstack((self.B, B)) if self.B.size else B
            self._b = np.hstack((self.b, b)) if self.b.size else b

        # TODO: handle duplicate constraints

        # SVD of A
        u, s, vh = np.linalg.svd(self.A, full_matrices=True)
        M, N = u.shape[0], vh.shape[1]
        K = min(M, N)
        rcond = np.finfo(self.A.dtype).eps * max(M, N)

        # Z = null space of A
        if self.constraints:
            tol = np.amax(s) * rcond
            large = s > tol
            num = np.sum(large, dtype=int)
            uk = u[:, :K]
            vhk = vh[:K, :]
            s = np.divide(1, s, where=large, out=s)
            s[(~large,)] = 0
            self._Ainv = np.matmul(vhk.T, np.multiply(s[..., np.newaxis], uk.T))
            self._y0 = np.dot(self.Ainv, self.b)
            self._Z = vh[num:, :].T.conj()
            self._dydx = np.dot(
                np.linalg.pinv(np.eye(self.dim_y) - np.dot(self.Ainv, self.B)), self.Z
            )
            self._dim_x = self.Z.shape[1]
        else:
            self._Ainv = np.array([[]])
            self._y0 = np.zeros((self.dim_y,))
            self._Z = np.eye(self.dim_y)
            self._dydx = self.Z
            self._dim_x = self.dim_y
<commit_msg>can't reference dimension getter methods before built<commit_after>import numpy as np
from abc import ABC, abstractmethod
from inspect import getfullargspec

from desc import config
from desc.backend import use_jax, jnp, jit
from desc.utils import Timer
from desc.io import IOAble
from desc.derivatives import Derivative
from desc.compute import arg_order

# XXX: could use `indicies` instead of `arg_order` in ObjectiveFunction loops


class ObjectiveFunction(IOAble):
    """Objective function comprised of one or more Objectives."""

    _io_attrs_ = ["objectives", "constraints"]

    def __init__(self, objectives, constraints=(), eq=None, use_jit=True):
        """Initialize an Objective Function.

        Parameters
        ----------
        objectives : tuple of Objective
            List of objectives to be targeted during optimization.
        constraints : tuple of Objective, optional
            List of objectives to be used as constraints during optimization.
        eq : Equilibrium, optional
            Equilibrium that will be optimized to satisfy the Objective.
        use_jit : bool, optional
            Whether to just-in-time compile the objective and derivatives.

        """
        if not isinstance(objectives, tuple):
            objectives = (objectives,)
        if not isinstance(constraints, tuple):
            constraints = (constraints,)

        self._objectives = objectives
        self._constraints = constraints
        self._use_jit = use_jit
        self._built = False
        self._compiled = False

        if eq is not None:
            self.build(eq, use_jit=self._use_jit)

    def _set_state_vector(self):
        """Set state vector components, dimensions, and indicies."""
        self._args = np.concatenate([obj.args for obj in self.objectives])
        if self.constraints:
            self._args = np.concatenate(
                (self.args, np.concatenate([obj.args for obj in self.constraints]))
            )
        self._args = np.unique(self.args)

        self._dimensions = self.objectives[0].dimensions

        idx = 0
        self._indicies = {}
        for arg in arg_order:
            if arg in self.args:
                self._indicies[arg] = np.arange(idx, idx + self.dimensions[arg])
                idx += self.dimensions[arg]
            else:
                self._indicies[arg] = np.array([])

        self._dim_y = idx

    def _build_linear_constraints(self):
        """Compute and factorize A to get pseudoinverse and nullspace."""
        # A*y = b = B*y
        self._A = np.array([[]])
        self._B = np.array([[]])
        self._b = np.array([])
        for obj in self.constraints:
            A = np.array([[]])
            for arg in arg_order:
                if arg in self.args:
                    a = np.atleast_2d(obj.derivatives[arg])
                    A = np.hstack((A, a)) if A.size else a
            if obj.target_arg in self.args:
                B = np.eye(self._dim_y)[self.indicies[obj.target_arg], :]
            else:
                B = np.zeros((obj._dim_f, self._dim_y))
            b = obj.target
            self._A = np.vstack((self.A, A)) if self.A.size else A
            self._B = np.vstack((self.B, B)) if self.B.size else B
            self._b = np.hstack((self.b, b)) if self.b.size else b

        # TODO: handle duplicate constraints

        # SVD of A
        u, s, vh = np.linalg.svd(self.A, full_matrices=True)
        M, N = u.shape[0], vh.shape[1]
        K = min(M, N)
        rcond = np.finfo(self.A.dtype).eps * max(M, N)

        # Z = null space of A
        if self.constraints:
            tol = np.amax(s) * rcond
            large = s > tol
            num = np.sum(large, dtype=int)
            uk = u[:, :K]
            vhk = vh[:K, :]
            s = np.divide(1, s, where=large, out=s)
            s[(~large,)] = 0
            self._Ainv = np.matmul(vhk.T, np.multiply(s[..., np.newaxis], uk.T))
            self._y0 = np.dot(self.Ainv, self.b)
            self._Z = vh[num:, :].T.conj()
            self._dydx = np.dot(
                np.linalg.pinv(np.eye(self._dim_y) - np.dot(self.Ainv, self.B)), self.Z
            )
            self._dim_x = self.Z.shape[1]
        else:
            self._Ainv = np.array([[]])
            self._y0 = np.zeros((self._dim_y,))
            self._Z = np.eye(self._dim_y)
            self._dydx = self.Z
            self._dim_x = self._dim_y

{'fraction_non_alphanumeric': 0.11494383297401566, 'fraction_numerical': 0.0022693747872461136, 'mean_word_length': 1.8259057390189164, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, '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': '4131017', 'n_tokens_mistral': 2758, 'n_tokens_neox': 2680, 'n_words': 726}