jacobian.py 文件源码

def _set_abs(self, abs_key, subjac):
        """
        Set sub-Jacobian.

        Parameters
        ----------
        abs_key : (str, str)
            Absolute name pair of sub-Jacobian.
        subjac : int or float or ndarray or sparse matrix
            sub-Jacobian as a scalar, vector, array, or AIJ list or tuple.
        """
        if not issparse(subjac):
            # np.promote_types will choose the smallest dtype that can contain both arguments
            subjac = np.atleast_1d(subjac)
            safe_dtype = np.promote_types(subjac.dtype, float)
            subjac = subjac.astype(safe_dtype, copy=False)

            # Bail here so that we allow top level jacobians to be of reduced size when indices are
            # specified on driver vars.
            if self._override_checks:
                self._subjacs[abs_key] = subjac
                return

            if abs_key in self._subjacs_info:
                subjac_info = self._subjacs_info[abs_key][0]
                rows = subjac_info['rows']
            else:
                rows = None

            if rows is None:
                # Dense subjac
                shape = self._abs_key2shape(abs_key)
                subjac = np.atleast_2d(subjac)
                if subjac.shape == (1, 1):
                    subjac = subjac[0, 0] * np.ones(shape, dtype=safe_dtype)
                else:
                    subjac = subjac.reshape(shape)

                if abs_key in self._subjacs and self._subjacs[abs_key].shape == shape:
                    np.copyto(self._subjacs[abs_key], subjac)
                else:
                    self._subjacs[abs_key] = subjac.copy()
            else:
                # Sparse subjac
                if subjac.shape == (1,):
                    subjac = subjac[0] * np.ones(rows.shape, dtype=safe_dtype)

                if subjac.shape != rows.shape:
                    raise ValueError("Sub-jacobian for key %s has "
                                     "the wrong shape (%s), expected (%s)." %
                                     (abs_key, subjac.shape, rows.shape))

                if abs_key in self._subjacs and subjac.shape == self._subjacs[abs_key][0].shape:
                    np.copyto(self._subjacs[abs_key][0], subjac)
                else:
                    self._subjacs[abs_key] = [subjac.copy(), rows, subjac_info['cols']]
        else:
            self._subjacs[abs_key] = subjac