jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def get_solution(self, area):
        '''
        Solve a truss problem

        Args:
            area : torch.Tensor, size = (N, n_bar), dtype = float64
                area of bars

        Returns:
            u : torch.Tensor, size=(n_node, n_dim), dtype=float64
                displacement of nodes
            tension : torch.Tensor, size=(n_bar, ), dtype=float64
                tension of bars
            volume : torch.Tensor, size=(n_bar, ), dtype=float64
                volume of bars
        '''

        N = area.shape[0]
        sys = torch.zeros((N, self.n_dim * self.n_node, self.n_dim * self.n_node), dtype=torch.float, device=self.device)  # system matrix
        volume = self.length * area  # volume of bars
        G = self.modulus * area / self.length  # a temporary coefficient to be used later

        local_matrix = G[:, :, None, None] * self.matrix  # local stiffness matrix
        for i in range(self.n_bar):
            sys[:, self.idx[i].unsqueeze(0), self.idx[i].unsqueeze(-1)] += local_matrix[:, i]  # global stiffness matrix

        u = torch.zeros((N, self.n_node, self.n_dim), dtype=torch.float, device=self.device)  # displacement

        u_free = torch.linalg.solve(sys[:, self.free_list][:, :, self.free_list], self.load[self.free_mask].unsqueeze(-1))
        u[:, self.free_mask] = u_free.squeeze()

        TG = torch.einsum('ij,jk->ijk', G, self.T)  # projection coefficient
        tension = torch.sum(TG * (u[:, self.con[:, 1], :] - u[:, self.con[:, 0], :]), dim=-1)
        return u, tension, volume
//  [python]
def se3_log_map(transform: torch.Tensor, eps: float = 1e-4, cos_bound: float = 1e-4, backend=None,
                test_acc=True):
    if backend is None:
        loguru.logger.warning("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
        loguru.logger.warning("!!!!se3_log_map backend is None!!!!")
        loguru.logger.warning("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
        backend = 'pytorch3d'
    if backend == 'pytorch3d':
        dof6 = pytorch3d.transforms.se3.se3_log_map(transform, eps, cos_bound)
    elif backend == 'opencv':
        from .pytorch3d_se3 import _se3_V_matrix, _get_se3_V_input
        # from pytorch3d.common.compat import solve
        log_rotation = []
        for tsfm in transform:
            cv2_rot = -cv2.Rodrigues(to_array(tsfm[:3, :3]))[0]
            log_rotation.append(torch.from_numpy(cv2_rot.reshape(-1)).to(transform.device).float())
        log_rotation = torch.stack(log_rotation, dim=0)
        T = transform[:, 3, :3]
        V = _se3_V_matrix(*_get_se3_V_input(log_rotation), eps=eps)
        log_translation = torch.solve(T[:, :, None], V)[0][:, :, 0]
        dof6 = torch.cat((log_translation, log_rotation), dim=1)
    else:
        raise NotImplementedError()
    if test_acc:
        err = (se3_exp_map(dof6) - transform).abs().max()
        if err > 0.1:
            raise RuntimeError()
    return dof6
//  [python]
def se3_log_map(transform: torch.Tensor, eps: float = 1e-4, cos_bound: float = 1e-4, backend=None,
                test_acc=True):
    if backend is None:
        loguru.logger.warning("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
        loguru.logger.warning("!!!!se3_log_map backend is None!!!!")
        loguru.logger.warning("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
        backend = 'pytorch3d'
    if backend == 'pytorch3d':
        dof6 = pytorch3d.transforms.se3.se3_log_map(transform, eps, cos_bound)
    elif backend == 'opencv':
        from .pytorch3d_se3 import _se3_V_matrix, _get_se3_V_input
        # from pytorch3d.common.compat import solve
        log_rotation = []
        for tsfm in transform:
            cv2_rot = -cv2.Rodrigues(to_array(tsfm[:3, :3]))[0]
            log_rotation.append(torch.from_numpy(cv2_rot.reshape(-1)).to(transform.device).float())
        log_rotation = torch.stack(log_rotation, dim=0)
        T = transform[:, 3, :3]
        V = _se3_V_matrix(*_get_se3_V_input(log_rotation), eps=eps)
        log_translation = torch.linalg.solve(V, T[:, :, None])[:, :, 0]
        dof6 = torch.cat((log_translation, log_rotation), dim=1)
    else:
        raise NotImplementedError()
    if test_acc:
        err = (se3_exp_map(dof6) - transform).abs().max()
        if err > 0.1:
            raise RuntimeError()
    return dof6
//  [python]
def helliger_distance_mvn(mean1, cov1, mean2, cov2):
	# Form (cov1 + cov2) / 2
	half_cov1_cov2 = (cov1 + cov2) / 2
	ret = torch.sqrt((torch.sqrt(torch.linalg.det(cov1)) * torch.sqrt(torch.linalg.det(cov2))) / torch.sqrt(torch.linalg.det(half_cov1_cov2)))
	ret *= torch.exp(-(1/8)*torch.matmul((mean1-mean2).t(), torch.solve(mean1-mean2, half_cov1_cov2)[0]))
	return 1.0 - ret
//  [python]
def helliger_distance_mvn(mean1, cov1, mean2, cov2):
	# Form (cov1 + cov2) / 2
	half_cov1_cov2 = (cov1 + cov2) / 2
	ret = torch.sqrt((torch.sqrt(torch.linalg.det(cov1)) * torch.sqrt(torch.linalg.det(cov2))) / torch.sqrt(torch.linalg.det(half_cov1_cov2)))
	ret *= torch.exp(-(1/8)*torch.matmul((mean1-mean2).t(), torch.linalg.solve(half_cov1_cov2, mean1-mean2)))
	return 1.0 - ret
//  [python]
def _log_gauss(self, x, mu, sigma):
        """
        :param x: 观测变量
        :param mu: 均值
        :param sigma: 协方差
        :return:
        """
        residual, data_dimension = x - mu, mu.shape[0]

        # sigma_inverse = torch.linalg.inv(sigma)
        log_exp_term = 0.5 * residual.transpose(1, 0) @ torch.solve(residual, sigma)[0]
        # print(f"[Info] The sigma inverse is {sigma_inverse}")
        previous_term = -1.0 * (0.5 * data_dimension * np.log(2 * np.pi) + 0.5 * torch.det(sigma))
        total_term = log_exp_term.squeeze() + previous_term.squeeze()
        return total_term
//  [python]
def _log_gauss(self, x, mu, sigma):
        """
        :param x: 观测变量
        :param mu: 均值
        :param sigma: 协方差
        :return:
        """
        residual, data_dimension = x - mu, mu.shape[0]

        # sigma_inverse = torch.linalg.inv(sigma)
        log_exp_term = 0.5 * residual.transpose(1, 0) @ torch.linalg.solve(sigma, residual)
        # print(f"[Info] The sigma inverse is {sigma_inverse}")
        previous_term = -1.0 * (0.5 * data_dimension * np.log(2 * np.pi) + 0.5 * torch.det(sigma))
        total_term = log_exp_term.squeeze() + previous_term.squeeze()
        return total_term
//  [python]
def forward(self, params: torch.Tensor) -> torch.Tensor:
        if self.orthogonal_map != "euler":
            # Construct a lower diagonal matrix where to place the parameters.
            offset = -1 if self.orthogonal_map == 'householder' else 0
            tril_indices = torch.tril_indices(row=self.d, col=self.d, offset=offset, device=params.device)
            new_params = torch.zeros(
                (params.size(0), self.d, self.d), dtype=params.dtype, device=params.device)
            new_params[:, tril_indices[0], tril_indices[1]] = params
            params = new_params

        if self.orthogonal_map == "matrix_exp" or self.orthogonal_map == "cayley":
            # We just need n x k - k(k-1)/2 parameters
            params = params.tril()
            A = params - params.transpose(-2, -1)
            # A is skew-symmetric (or skew-hermitian)
            if self.orthogonal_map == "matrix_exp":
                Q = torch.matrix_exp(A)
            elif self.orthogonal_map == "cayley":
                # Computes the Cayley retraction (I+A/2)(I-A/2)^{-1}
                Id = torch.eye(self.d, dtype=A.dtype, device=A.device)
                Q = torch.solve(torch.add(Id, A, alpha=0.5), torch.add(Id, A, alpha=-0.5))[1]
        elif self.orthogonal_map == 'householder':
            eye = torch.eye(self.d, device=params.device).unsqueeze(0).repeat(params.size(0), 1, 1)
            A = params.tril(diagonal=-1) + eye
            Q = torch_householder_orgqr(A)
        elif self.orthogonal_map == 'euler':
            assert 2 <= self.d <= 3
            if self.d == 2:
                Q = self.get_2d_rotation(params)
            else:
                Q = self.get_3d_rotation(params)
        else:
            raise ValueError(f"Unsupported transformations {self.orthogonal_map}")
        return Q
//  [python]
def forward(self, params: torch.Tensor) -> torch.Tensor:
        if self.orthogonal_map != "euler":
            # Construct a lower diagonal matrix where to place the parameters.
            offset = -1 if self.orthogonal_map == 'householder' else 0
            tril_indices = torch.tril_indices(row=self.d, col=self.d, offset=offset, device=params.device)
            new_params = torch.zeros(
                (params.size(0), self.d, self.d), dtype=params.dtype, device=params.device)
            new_params[:, tril_indices[0], tril_indices[1]] = params
            params = new_params

        if self.orthogonal_map == "matrix_exp" or self.orthogonal_map == "cayley":
            # We just need n x k - k(k-1)/2 parameters
            params = params.tril()
            A = params - params.transpose(-2, -1)
            # A is skew-symmetric (or skew-hermitian)
            if self.orthogonal_map == "matrix_exp":
                Q = torch.matrix_exp(A)
            elif self.orthogonal_map == "cayley":
                # Computes the Cayley retraction (I+A/2)(I-A/2)^{-1}
                Id = torch.eye(self.d, dtype=A.dtype, device=A.device)
                Q = torch.linalg.solve(torch.add(Id, A, alpha=-0.5), torch.add(Id, A, alpha=0.5))
        elif self.orthogonal_map == 'householder':
            eye = torch.eye(self.d, device=params.device).unsqueeze(0).repeat(params.size(0), 1, 1)
            A = params.tril(diagonal=-1) + eye
            Q = torch_householder_orgqr(A)
        elif self.orthogonal_map == 'euler':
            assert 2 <= self.d <= 3
            if self.d == 2:
                Q = self.get_2d_rotation(params)
            else:
                Q = self.get_3d_rotation(params)
        else:
            raise ValueError(f"Unsupported transformations {self.orthogonal_map}")
        return Q
//  [python]
def _batch_regress_yx(self, y, fx, y_prev, fx_prev):
        # for the first batch use the same data to regress on.
        if y_prev == None:
            y_prev = y.clone().detach()
        if fx_prev == None:
            fx_prev = fx.clone().detach()

        if self.model_cfg.regression == 'linear':
            y_prev = y_prev.cpu().numpy()
            fx_prev = fx_prev.clone().detach().cpu().numpy()
            linear_reg = linear_model.LinearRegression()
            linear_reg.fit(y_prev, fx_prev)

            fx_ = np.matmul(y.cpu().numpy(), linear_reg.coef_.T) + linear_reg.intercept_
            return torch.FloatTensor(fx_).to(self.device)

        elif self.model_cfg.regression == 'linear_old':
            y = y.cpu().numpy()
            fx = fx.clone().detach().cpu().numpy()
            linear_reg = linear_model.LinearRegression()
            linear_reg.fit(y, fx)

            fx_ = np.matmul(y, linear_reg.coef_.T) + linear_reg.intercept_
            return torch.FloatTensor(fx_).to(self.device)

        elif self.model_cfg.regression == 'kernel-ridge':
            KY = eval(f'losses.{self.kernel_y}_kernel(y_prev, **self.kernel_y_args)')
            I = torch.eye(KY.shape[0], device=KY.device)
            W = torch.solve(fx_prev, (KY + self.model_cfg.ridge_lambda*I))[0]

            KyY = eval(f'losses.{self.kernel_y}_kernel(y, Y=y_prev, **self.kernel_y_args)')
            fx_ = KyY @ W
            return fx_
//  [python]
def _batch_regress_yx(self, y, fx, y_prev, fx_prev):
        # for the first batch use the same data to regress on.
        if y_prev == None:
            y_prev = y.clone().detach()
        if fx_prev == None:
            fx_prev = fx.clone().detach()

        if self.model_cfg.regression == 'linear':
            y_prev = y_prev.cpu().numpy()
            fx_prev = fx_prev.clone().detach().cpu().numpy()
            linear_reg = linear_model.LinearRegression()
            linear_reg.fit(y_prev, fx_prev)

            fx_ = np.matmul(y.cpu().numpy(), linear_reg.coef_.T) + linear_reg.intercept_
            return torch.FloatTensor(fx_).to(self.device)

        elif self.model_cfg.regression == 'linear_old':
            y = y.cpu().numpy()
            fx = fx.clone().detach().cpu().numpy()
            linear_reg = linear_model.LinearRegression()
            linear_reg.fit(y, fx)

            fx_ = np.matmul(y, linear_reg.coef_.T) + linear_reg.intercept_
            return torch.FloatTensor(fx_).to(self.device)

        elif self.model_cfg.regression == 'kernel-ridge':
            KY = eval(f'losses.{self.kernel_y}_kernel(y_prev, **self.kernel_y_args)')
            I = torch.eye(KY.shape[0], device=KY.device)
            W = torch.linalg.solve((KY + self.model_cfg.ridge_lambda*I), fx_prev)

            KyY = eval(f'losses.{self.kernel_y}_kernel(y, Y=y_prev, **self.kernel_y_args)')
            fx_ = KyY @ W
            return fx_
//  [python]
def f(W, x):
    return torch.abs(W @ x) + torch.solve(x,W)[1]
//  [python]
def f(W, x):
    return torch.abs(W @ x) + torch.linalg.solve(W, x)
//  [python]
def fast_global_registration(xyz, xyz_target, mu_init=1, num_iter=20):
    """
    estimate the rotation and translation using Fast Global Registration algorithm (M estimator for robust estimation)
    http://vladlen.info/papers/fast-global-registration.pdf
    """
    assert xyz.shape == xyz_target.shape

    T_res = torch.eye(4, device=xyz.device)
    mu = mu_init
    source = xyz.clone()
    weight = torch.ones(len(source), 1, device=xyz.device)
    for i in range(num_iter):
        if i > 0 and i % 5 == 0:
            mu /= 2.0
        A, b = get_matrix_system(source, xyz_target, weight)
        x = torch.solve(A.T.mm(A), A.T @ b)[0]
        T = get_trans(x.view(-1))
        source = source.mm(T[:3, :3].T) + T[:3, 3]
        T_res = T @ T_res
        weight = get_geman_mclure_weight(source, xyz_target, mu)
    return T_res
//  [python]
def fast_global_registration(xyz, xyz_target, mu_init=1, num_iter=20):
    """
    estimate the rotation and translation using Fast Global Registration algorithm (M estimator for robust estimation)
    http://vladlen.info/papers/fast-global-registration.pdf
    """
    assert xyz.shape == xyz_target.shape

    T_res = torch.eye(4, device=xyz.device)
    mu = mu_init
    source = xyz.clone()
    weight = torch.ones(len(source), 1, device=xyz.device)
    for i in range(num_iter):
        if i > 0 and i % 5 == 0:
            mu /= 2.0
        A, b = get_matrix_system(source, xyz_target, weight)
        x = torch.linalg.solve(A.T.mm(A), A.T @ b)
        T = get_trans(x.view(-1))
        source = source.mm(T[:3, :3].T) + T[:3, 3]
        T_res = T @ T_res
        weight = get_geman_mclure_weight(source, xyz_target, mu)
    return T_res
//  [python]
def rts_smoother(self, filtering_mean, filtering_cov):

        temp_smoothing_mean, temp_smoothing_cov, gain_list = [], [], []
        temp_smoothing_mean.append(filtering_mean[-1])
        temp_smoothing_cov.append(filtering_cov[-1])
        # print(len(filtering_mean))

        for t in range(self.seq_length - 1, -1, -1):
            # print(f'[Info] Time {t} start smoothing!')
            # 先验均值、方差
            prior_mean = self.transition_matrix @ filtering_mean[t]
            prior_cov = self.transition_matrix @ filtering_cov[t] @ self.transition_matrix.transpose(1, 0) + self.transition_noise

            # rts gain
            residual_mean = temp_smoothing_mean[-1] - prior_mean
            residual_cov = temp_smoothing_cov[-1] - prior_cov
            gain_matrix = filtering_cov[t] @ self.transition_matrix.transpose(1, 0) @ torch.inverse(prior_cov)

            # 平滑均值、方差
            smooth_mean = filtering_mean[t] + filtering_cov[t] @ self.transition_matrix.transpose(1, 0) @ torch.solve(residual_mean, prior_cov)[0] #residual_mean
            smooth_cov = filtering_cov[t] + gain_matrix @ residual_cov @ gain_matrix.transpose(1, 0)

            temp_smoothing_mean.append(smooth_mean)
            temp_smoothing_cov.append(smooth_cov)
            gain_list.append(gain_matrix)

        temp_smoothing_mean = list(reversed(temp_smoothing_mean))
        temp_smoothing_cov = list(reversed(temp_smoothing_cov))
        return temp_smoothing_mean, temp_smoothing_cov, gain_list
//  [python]
def rts_smoother(self, filtering_mean, filtering_cov):

        temp_smoothing_mean, temp_smoothing_cov, gain_list = [], [], []
        temp_smoothing_mean.append(filtering_mean[-1])
        temp_smoothing_cov.append(filtering_cov[-1])
        # print(len(filtering_mean))

        for t in range(self.seq_length - 1, -1, -1):
            # print(f'[Info] Time {t} start smoothing!')
            # 先验均值、方差
            prior_mean = self.transition_matrix @ filtering_mean[t]
            prior_cov = self.transition_matrix @ filtering_cov[t] @ self.transition_matrix.transpose(1, 0) + self.transition_noise

            # rts gain
            residual_mean = temp_smoothing_mean[-1] - prior_mean
            residual_cov = temp_smoothing_cov[-1] - prior_cov
            gain_matrix = filtering_cov[t] @ self.transition_matrix.transpose(1, 0) @ torch.linalg.inv(prior_cov)

            # 平滑均值、方差
            smooth_mean = filtering_mean[t] + filtering_cov[t] @ self.transition_matrix.transpose(1, 0) @ torch.linalg.solve(prior_cov, residual_mean) #residual_mean
            smooth_cov = filtering_cov[t] + gain_matrix @ residual_cov @ gain_matrix.transpose(1, 0)

            temp_smoothing_mean.append(smooth_mean)
            temp_smoothing_cov.append(smooth_cov)
            gain_list.append(gain_matrix)

        temp_smoothing_mean = list(reversed(temp_smoothing_mean))
        temp_smoothing_cov = list(reversed(temp_smoothing_cov))
        return temp_smoothing_mean, temp_smoothing_cov, gain_list
//  [python]
def _train_mle(self, x, y):

        reg_coeff = 1e-6
        alpha_mle = torch.solve((x.T @ x + reg_coeff * torch.eye(self.nb_coeff).to(x)), x.T @ y)[0]

        return alpha_mle
//  [python]
def _train_mle(self, x, y):

        reg_coeff = 1e-6
        alpha_mle = torch.linalg.solve((x.T @ x + reg_coeff * torch.eye(self.nb_coeff).to(x)), x.T @ y)

        return alpha_mle
//  [python]
def kalman_filter(self, data):
        """
        :param data:
        :return:
        """
        temp_filtering_mean, temp_filtering_cov = [], []
        likelihood_list = []
        temp_filtering_mean.append(self.init_mean)
        temp_filtering_cov.append(self.init_cov)


        for t in range(self.seq_length):

            # 先验转移
            prior_mean = self.transition_matrix @ temp_filtering_mean[-1]
            prior_cov = self.transition_matrix @ temp_filtering_cov[-1] @ self.transition_matrix.transpose(1, 0) + self.transition_noise

            # 观测方程
            obs_mean = self.observation_matrix @ prior_mean
            obs_cov = self.observation_matrix @ prior_cov @ self.observation_matrix.transpose(1, 0) + self.observation_noise

            # 新息
            residual = data[:, t:t+1] - obs_mean
            # kalman_gain = prior_cov @ self.observation_matrix.transpose(1, 0) @ torch.linalg.inv(obs_cov)
            # kalman_gain = prior_cov @ self.observation_matrix.transpose(1, 0) @ self._self_inverse(obs_cov)

            # 更新滤波均值、滤波方差
            filtered_mean = prior_mean + prior_cov @ self.observation_matrix.transpose(1, 0) @ torch.solve(residual, obs_cov)[0]
            # filtered_cov = prior_cov - kalman_gain @ obs_cov @ kalman_gain.transpose(1, 0)
            filtered_cov = prior_cov - prior_cov @ self.observation_matrix.transpose(1, 0) @ torch.inverse(obs_cov) @ self.observation_matrix @ prior_cov
            temp_filtering_mean.append(filtered_mean)
            temp_filtering_cov.append(filtered_cov)

            likelihood_list.append(self._log_gauss(x=data[:, t:t+1],
                                                   mu=obs_mean,
                                                   sigma=obs_cov))

        return temp_filtering_mean, temp_filtering_cov, likelihood_list
//  [python]
def kalman_filter(self, data):
        """
        :param data:
        :return:
        """
        temp_filtering_mean, temp_filtering_cov = [], []
        likelihood_list = []
        temp_filtering_mean.append(self.init_mean)
        temp_filtering_cov.append(self.init_cov)


        for t in range(self.seq_length):

            # 先验转移
            prior_mean = self.transition_matrix @ temp_filtering_mean[-1]
            prior_cov = self.transition_matrix @ temp_filtering_cov[-1] @ self.transition_matrix.transpose(1, 0) + self.transition_noise

            # 观测方程
            obs_mean = self.observation_matrix @ prior_mean
            obs_cov = self.observation_matrix @ prior_cov @ self.observation_matrix.transpose(1, 0) + self.observation_noise

            # 新息
            residual = data[:, t:t+1] - obs_mean
            # kalman_gain = prior_cov @ self.observation_matrix.transpose(1, 0) @ torch.linalg.inv(obs_cov)
            # kalman_gain = prior_cov @ self.observation_matrix.transpose(1, 0) @ self._self_inverse(obs_cov)

            # 更新滤波均值、滤波方差
            filtered_mean = prior_mean + prior_cov @ self.observation_matrix.transpose(1, 0) @ torch.linalg.solve(obs_cov, residual)
            # filtered_cov = prior_cov - kalman_gain @ obs_cov @ kalman_gain.transpose(1, 0)
            filtered_cov = prior_cov - prior_cov @ self.observation_matrix.transpose(1, 0) @ torch.linalg.inv(obs_cov) @ self.observation_matrix @ prior_cov
            temp_filtering_mean.append(filtered_mean)
            temp_filtering_cov.append(filtered_cov)

            likelihood_list.append(self._log_gauss(x=data[:, t:t+1],
                                                   mu=obs_mean,
                                                   sigma=obs_cov))

        return temp_filtering_mean, temp_filtering_cov, likelihood_list
//  [python]
def forward(self, X):
        # Assume X skew-symmetric
        # (I + X)(I - X)^{-1}
        return torch.solve(self.Id + X, self.Id - X)[1]
//  [python]
def forward(self, X):
        # (I + X)(I - X)^{-1}
        return torch.solve(self.Id + X, self.Id - X)[1]
//  [python]
def fastmvg(Phi, alpha, D):
    # Fast sampler for multivariate Gaussian distributions (large p, p > n) of
    #  the form N(mu, S), where
    #       mu = S Phi' y
    #       S  = inv(Phi'Phi + inv(D))
    # Reference:
    #   Fast sampling with Gaussian scale-mixture priors in high-dimensional
    #   regression, A. Bhattacharya, A. Chakraborty and B. K. Mallick
    #   arXiv:1506.04778

    n, p = Phi.shape

    d = D.diag().to(Phi)
    u = torch.randn(p).to(Phi) * d.sqrt()
    delta = torch.randn(n).to(Phi)
    v = (Phi @ u + delta).view(-1, 1)
    Dpt = Phi.T * d[:, None]
    w = torch.solve(alpha - v, Phi @ Dpt + torch.eye(n).to(Phi))[0]
    x = u + (Dpt @ w).view(-1)

    return x
//  [python]
def fastmvg(Phi, alpha, D):
    # Fast sampler for multivariate Gaussian distributions (large p, p > n) of
    #  the form N(mu, S), where
    #       mu = S Phi' y
    #       S  = inv(Phi'Phi + inv(D))
    # Reference:
    #   Fast sampling with Gaussian scale-mixture priors in high-dimensional
    #   regression, A. Bhattacharya, A. Chakraborty and B. K. Mallick
    #   arXiv:1506.04778

    n, p = Phi.shape

    d = D.diag().to(Phi)
    u = torch.randn(p).to(Phi) * d.sqrt()
    delta = torch.randn(n).to(Phi)
    v = (Phi @ u + delta).view(-1, 1)
    Dpt = Phi.T * d[:, None]
    w = torch.linalg.solve(Phi @ Dpt + torch.eye(n).to(Phi), alpha - v)
    x = u + (Dpt @ w).view(-1)

    return x
//  [python]
def _train_bayes(self, x, y, a_prior, b_prior):

        n = x.shape[0]

        # set alpha prior parameters
        alpha_mean_prior = (np.pi / 2) / self.nb_coeff * torch.ones((self.nb_coeff, 1)).to(x)
        alpha_cov_prior = torch.eye(self.nb_coeff).to(x)
        alpha_inv_prior = torch.inverse(alpha_cov_prior)

        # compute alpha posterior parameters
        alpha_prec_ = alpha_inv_prior + x.T @ x
        temp = torch.solve(x @ alpha_cov_prior, torch.eye(n).to(x) + x @ alpha_cov_prior @ x.T)[0]
        alpha_cov = alpha_cov_prior - alpha_cov_prior @ x.T @ temp  # Sherman-Morrison formula
        alpha_cov = (alpha_cov + alpha_cov.T) / 2  # symmetrizing covariance
        alpha_mean = alpha_cov @ (alpha_inv_prior @ alpha_mean_prior + x.T @ y)

        # compute hyper-parameter posterior parameters. Equation 88 and 89 https://www.cs.ubc.ca/~murphyk/Papers/bayesGauss.pdf
        a_post = a_prior + n / 2
        b_post = b_prior + 0.5 * (
                alpha_mean_prior.T @ alpha_inv_prior @ alpha_mean_prior + y.T @ y - alpha_mean.T @ alpha_prec_ @ alpha_mean).item()

        return alpha_mean, alpha_cov, a_post, b_post
//  [python]
def _train_bayes(self, x, y, a_prior, b_prior):

        n = x.shape[0]

        # set alpha prior parameters
        alpha_mean_prior = (np.pi / 2) / self.nb_coeff * torch.ones((self.nb_coeff, 1)).to(x)
        alpha_cov_prior = torch.eye(self.nb_coeff).to(x)
        alpha_inv_prior = torch.linalg.inv(alpha_cov_prior)

        # compute alpha posterior parameters
        alpha_prec_ = alpha_inv_prior + x.T @ x
        temp = torch.linalg.solve(torch.eye(n).to(x) + x @ alpha_cov_prior @ x.T, x @ alpha_cov_prior)
        alpha_cov = alpha_cov_prior - alpha_cov_prior @ x.T @ temp  # Sherman-Morrison formula
        alpha_cov = (alpha_cov + alpha_cov.T) / 2  # symmetrizing covariance
        alpha_mean = alpha_cov @ (alpha_inv_prior @ alpha_mean_prior + x.T @ y)

        # compute hyper-parameter posterior parameters. Equation 88 and 89 https://www.cs.ubc.ca/~murphyk/Papers/bayesGauss.pdf
        a_post = a_prior + n / 2
        b_post = b_prior + 0.5 * (
                alpha_mean_prior.T @ alpha_inv_prior @ alpha_mean_prior + y.T @ y - alpha_mean.T @ alpha_prec_ @ alpha_mean).item()

        return alpha_mean, alpha_cov, a_post, b_post
//  [python]
def right_inverse(self, A):
        # Assume A orthogonal
        # See https://en.wikipedia.org/wiki/Cayley_transform#Matrix_map
        # (A - I)(A + I)^{-1}
        return torch.solve(self.Id - A, A + self.Id)[1]
//  [python]
def right_inverse(self, A):
        # Assume A orthogonal
        # See https://en.wikipedia.org/wiki/Cayley_transform#Matrix_map
        # (A - I)(A + I)^{-1}
        return torch.linalg.solve(A + self.Id, self.Id - A)
//  [python]
def fastmvg_rue(X, XtX, y, Sigma_star):
    # Another sampler for multivariate Gaussians (small p) of the form
    #  N(mu, S), where
    #  mu = S Phi' y
    #  S  = inv(Phi'Phi + inv(D))
    #
    # Here, PtP = Phi'*Phi (X'X is precomputed)
    #

    p = X.shape[1]
    Sigma_star_inv = (1. / Sigma_star.diag()).diag()

    # regularize PtP + Dinv matrix for small negative eigenvalues
    try:
        L = torch.linalg.cholesky(XtX + Sigma_star_inv)
    except:
        mat = XtX + Sigma_star_inv
        Smat = (mat + mat.T) / 2.
        maxEig_Smat = torch.max(torch.linalg.eigvals(Smat).real)
        L = torch.linalg.cholesky(Smat + maxEig_Smat * 1e-15 * torch.eye(Smat.shape[0]).to(X))

    v = torch.solve(X.T @ y, L)[0]
    m = torch.solve(v, L.T)[0]
    w = torch.solve(torch.randn((p, 1)).to(X), L.T)[0]

    x = m + w

    return x.view(-1)
//  [python]
def fastmvg_rue(X, XtX, y, Sigma_star):
    # Another sampler for multivariate Gaussians (small p) of the form
    #  N(mu, S), where
    #  mu = S Phi' y
    #  S  = inv(Phi'Phi + inv(D))
    #
    # Here, PtP = Phi'*Phi (X'X is precomputed)
    #

    p = X.shape[1]
    Sigma_star_inv = (1. / Sigma_star.diag()).diag()

    # regularize PtP + Dinv matrix for small negative eigenvalues
    try:
        L = torch.linalg.cholesky(XtX + Sigma_star_inv)
    except:
        mat = XtX + Sigma_star_inv
        Smat = (mat + mat.T) / 2.
        maxEig_Smat = torch.max(torch.linalg.eigvals(Smat).real)
        L = torch.linalg.cholesky(Smat + maxEig_Smat * 1e-15 * torch.eye(Smat.shape[0]).to(X))

    v = torch.linalg.solve(L, X.T @ y)
    m = torch.linalg.solve(L.T, v)
    w = torch.linalg.solve(L.T, torch.randn((p, 1)).to(X))

    x = m + w

    return x.view(-1)
//  [python]
def test_bicgstab(self):
        size = 100
        matrix_dense = torch.randn(size, size, dtype=torch.float64, device=self.device) + 10 * torch.eye(
            size, device=self.device
        )
        matrix_sparse = matrix_dense.to_sparse_csr()

        rhs = torch.randn(size, dtype=torch.float64, device=self.device)

        actual = torch.solve(matrix_dense, rhs)[0]
        # print("actual",actual)

        # mn = matrix_dense.cpu().numpy()
        # rn = rhs.cpu().numpy()
        # solves = scipy.sparse.linalg.bicgstab(mn,rn)
        # print("np solves",solves)

        solves = bicgstab(matrix_dense, rhs=rhs)
        # print("our solves",solves)
        self.assertTrue(torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))

        solves = bicgstab(matrix_dense.matmul, rhs=rhs)
        self.assertTrue(torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))

        solves = bicgstab(matrix_sparse.matmul, rhs=rhs)
        self.assertTrue(torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))

        solves = bicgstab(matrix_sparse, rhs=rhs)
        self.assertTrue(torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))
//  [python]
def test_bicgstab(self):
        size = 100
        matrix_dense = torch.randn(size, size, dtype=torch.float64, device=self.device) + 10 * torch.eye(
            size, device=self.device
        )
        matrix_sparse = matrix_dense.to_sparse_csr()

        rhs = torch.randn(size, dtype=torch.float64, device=self.device)

        actual = torch.linalg.solve(matrix_dense, rhs)
        # print("actual",actual)

        # mn = matrix_dense.cpu().numpy()
        # rn = rhs.cpu().numpy()
        # solves = scipy.sparse.linalg.bicgstab(mn,rn)
        # print("np solves",solves)

        solves = bicgstab(matrix_dense, rhs=rhs)
        # print("our solves",solves)
        self.assertTrue(torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))

        solves = bicgstab(matrix_dense.matmul, rhs=rhs)
        self.assertTrue(torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))

        solves = bicgstab(matrix_sparse.matmul, rhs=rhs)
        self.assertTrue(torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))

        solves = bicgstab(matrix_sparse, rhs=rhs)
        self.assertTrue(torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))
//  [python]
def inverse_mult(self, u, delta, transpose=False):
        """Computes (I - d A)^-1 u."""

        if isinstance(delta, torch.Tensor):
            delta = delta.unsqueeze(-1).unsqueeze(-1)
        _A = self.I - delta * self.A
        if transpose: _A = _A.transpose(-1, -2)

        # x = torch.linalg.solve(_A, u.unsqueeze(-1)).squeeze(-1)

        # TODO pass in a flag to toggle the two codepaths depending on how big the problem is
        xs = []
        for _A_, u_ in zip(*torch.broadcast_tensors(_A, u.unsqueeze(-1))):
            x_ = torch.solve(_A_, u_[...,:1])[0].squeeze(-1)
            xs.append(x_)
        x = torch.stack(xs, dim=0)

        return x
//  [python]
def inverse_mult(self, u, delta, transpose=False):
        """Computes (I - d A)^-1 u."""

        if isinstance(delta, torch.Tensor):
            delta = delta.unsqueeze(-1).unsqueeze(-1)
        _A = self.I - delta * self.A
        if transpose: _A = _A.transpose(-1, -2)

        # x = torch.linalg.solve(_A, u.unsqueeze(-1)).squeeze(-1)

        # TODO pass in a flag to toggle the two codepaths depending on how big the problem is
        xs = []
        for _A_, u_ in zip(*torch.broadcast_tensors(_A, u.unsqueeze(-1))):
            x_ = torch.linalg.solve(_A_, u_[...,:1]).squeeze(-1)
            xs.append(x_)
        x = torch.stack(xs, dim=0)

        return x
//  [python]
def _setup_linear(self):
        """Preprocessing that allows fast linear-time (in state dimension) stepping."""
        dt, A, B, C, P, Q = self._get_params()

        # Prepare Linear stepping
        D = (2.0 / dt - A).reciprocal()  # (H, N)
        R = (torch.eye(self.rank, dtype=A.dtype, device=A.device) + 2*contract('r h n, h n, s h n -> h r s', Q, D, P).real) # (H R R)
        Q_D = rearrange(Q*D, 'r h n -> h r n')
        try:
            R = torch.solve(Q_D, R)[0] # (H R N)
        except:
            R = torch.tensor(np.linalg.solve(R.to(Q_D).contiguous().detach().cpu(), Q_D.contiguous().detach().cpu())).to(Q_D)
        R = rearrange(R, 'h r n -> r h n')

        self.step_params = {
            "D": D, # (H N)
            "R": R, # (R H N)
            "P": P, # (R H N)
            "Q": Q, # (R H N)
            "B": B, # (1 H N)
            "E": 2.0 / dt + A, # (H N)
        }
//  [python]
def _setup_linear(self):
        """Preprocessing that allows fast linear-time (in state dimension) stepping."""
        dt, A, B, C, P, Q = self._get_params()

        # Prepare Linear stepping
        D = (2.0 / dt - A).reciprocal()  # (H, N)
        R = (torch.eye(self.rank, dtype=A.dtype, device=A.device) + 2*contract('r h n, h n, s h n -> h r s', Q, D, P).real) # (H R R)
        Q_D = rearrange(Q*D, 'r h n -> h r n')
        try:
            R = torch.linalg.solve(R, Q_D) # (H R N)
        except:
            R = torch.tensor(np.linalg.solve(R.to(Q_D).contiguous().detach().cpu(), Q_D.contiguous().detach().cpu())).to(Q_D)
        R = rearrange(R, 'h r n -> r h n')

        self.step_params = {
            "D": D, # (H N)
            "R": R, # (R H N)
            "P": P, # (R H N)
            "Q": Q, # (R H N)
            "B": B, # (1 H N)
            "E": 2.0 / dt + A, # (H N)
        }
//  [python]
def inverse_mult(self, u, delta, transpose=False):
        """ Computes (I - d A)^-1 u """

        if isinstance(delta, torch.Tensor):
            delta = delta.unsqueeze(-1).unsqueeze(-1)
        _A = self.I - delta * self.A
        if transpose: _A = _A.transpose(-1, -2)

        # x = torch.linalg.solve(_A, u.unsqueeze(-1)).squeeze(-1) # this can run out of memory
        xs = []
        for _A_, u_ in zip(*torch.broadcast_tensors(_A, u.unsqueeze(-1))):
            x_ = torch.solve(_A_, u_[...,:1])[0].squeeze(-1)
            xs.append(x_)
        x = torch.stack(xs, dim=0)

        return x
//  [python]
def inverse_mult(self, u, delta, transpose=False):
        """ Computes (I - d A)^-1 u """

        if isinstance(delta, torch.Tensor):
            delta = delta.unsqueeze(-1).unsqueeze(-1)
        _A = self.I - delta * self.A
        if transpose: _A = _A.transpose(-1, -2)

        # x = torch.linalg.solve(_A, u.unsqueeze(-1)).squeeze(-1) # this can run out of memory
        xs = []
        for _A_, u_ in zip(*torch.broadcast_tensors(_A, u.unsqueeze(-1))):
            x_ = torch.linalg.solve(_A_, u_[...,:1]).squeeze(-1)
            xs.append(x_)
        x = torch.stack(xs, dim=0)

        return x
//  [python]
def solve_adjoint(self, fwd_sol, param_vals, **newton_kwargs):
        jac = self._fwd_solver.physics._raw_residual(fwd_sol)[1]
        if jac is None:
            fwd_sol.requires_grad_(True)
            assert fwd_sol.requires_grad
            jac = torch.autograd.functional.jacobian(
                lambda s: self.physics._raw_residual(s)[0], fwd_sol,
                strict=True)
            fwd_sol.requires_grad_(False)
        fwd_sol_copy = torch.clone(fwd_sol)
        param_vals_copy = torch.clone(param_vals)
        if self._dqdu is None:
            functional_vals = self._functional(
                fwd_sol_copy.requires_grad_(True),
                param_vals_copy.requires_grad_(True))
            functional_vals.backward()
            dqdu = fwd_sol_copy.grad
        else:
            functional_vals = self._functional(fwd_sol_copy, param_vals_copy)
            dqdu = self._dqdu(fwd_sol_copy, param_vals_copy)
        __, tmp = (
            self._fwd_solver.physics.mesh._apply_boundary_conditions(
                self._fwd_solver.physics._bndry_conds, -dqdu, jac.clone(),
                fwd_sol, self._fwd_solver.physics.flux_jac))
        jac_adjoint = tmp.T
        rhs = -dqdu
        adj_sol = torch.solve(rhs, jac_adjoint)[0]
        return adj_sol
//  [python]
def solve_adjoint(self, fwd_sol, param_vals, **newton_kwargs):
        jac = self._fwd_solver.physics._raw_residual(fwd_sol)[1]
        if jac is None:
            fwd_sol.requires_grad_(True)
            assert fwd_sol.requires_grad
            jac = torch.autograd.functional.jacobian(
                lambda s: self.physics._raw_residual(s)[0], fwd_sol,
                strict=True)
            fwd_sol.requires_grad_(False)
        fwd_sol_copy = torch.clone(fwd_sol)
        param_vals_copy = torch.clone(param_vals)
        if self._dqdu is None:
            functional_vals = self._functional(
                fwd_sol_copy.requires_grad_(True),
                param_vals_copy.requires_grad_(True))
            functional_vals.backward()
            dqdu = fwd_sol_copy.grad
        else:
            functional_vals = self._functional(fwd_sol_copy, param_vals_copy)
            dqdu = self._dqdu(fwd_sol_copy, param_vals_copy)
        __, tmp = (
            self._fwd_solver.physics.mesh._apply_boundary_conditions(
                self._fwd_solver.physics._bndry_conds, -dqdu, jac.clone(),
                fwd_sol, self._fwd_solver.physics.flux_jac))
        jac_adjoint = tmp.T
        rhs = -dqdu
        adj_sol = torch.linalg.solve(jac_adjoint, rhs)
        return adj_sol
//  [python]
def get_polynomial(yy):
        '''
        Calculate the coefficients of cubic splines.
        Args:
            yy: [batchsize, n_sample]: the coarse {wi} on each ray
        Return: 
            coeff: [batchsize, n_sample-1, 4]: the coefficients of cubic splines of each interval on each ray
            degree: a number: the degree of splines
        '''
        degree = 3
        bins = torch.linspace(0, yy.shape[-1]-1, yy.shape[-1], device=yy.device).expand(yy.shape)
        batch_size = bins.size()[0]
        pts_num = bins.size()[1]
        bin_length = bins[..., 1:] - bins[..., :-1] # [b, n]
        matrix = torch.diag_embed(4 * torch.ones_like(bin_length[..., :-1]))
        for i in range(pts_num - 3):
            matrix[..., i, i + 1] = 1
            matrix[..., i + 1, i] = 1
        bi = 6 * (yy[..., 1:] - yy[..., :-1]).contiguous() # [b, n]
        vi = bi[..., 1:] - bi[..., :-1]
        second_dev = torch.solve(vi, matrix)[0]# [b, n-1]
        second_dev = torch.cat([torch.zeros_like(second_dev[..., :1]), second_dev, torch.zeros_like(second_dev[..., :1])], dim = -1)# [b, n+1]

        coeff = torch.zeros([batch_size, pts_num - 1, 4], device=bins.device)
        coeff[..., 3] = (second_dev[..., 1:] - second_dev[..., :-1]) / (6 * bin_length)
        coeff[..., 2] = second_dev[..., :-1] / 2
        coeff[..., 1] = -bin_length*second_dev[..., 1:]/6 - bin_length*second_dev[..., :-1]/3 + (yy[..., 1:] - yy[..., :-1]) / bin_length
        coeff[..., 0] = yy[..., :-1]
    
        return coeff, degree
//  [python]
def get_polynomial(yy):
        '''
        Calculate the coefficients of cubic splines.
        Args:
            yy: [batchsize, n_sample]: the coarse {wi} on each ray
        Return: 
            coeff: [batchsize, n_sample-1, 4]: the coefficients of cubic splines of each interval on each ray
            degree: a number: the degree of splines
        '''
        degree = 3
        bins = torch.linspace(0, yy.shape[-1]-1, yy.shape[-1], device=yy.device).expand(yy.shape)
        batch_size = bins.size()[0]
        pts_num = bins.size()[1]
        bin_length = bins[..., 1:] - bins[..., :-1] # [b, n]
        matrix = torch.diag_embed(4 * torch.ones_like(bin_length[..., :-1]))
        for i in range(pts_num - 3):
            matrix[..., i, i + 1] = 1
            matrix[..., i + 1, i] = 1
        bi = 6 * (yy[..., 1:] - yy[..., :-1]).contiguous() # [b, n]
        vi = bi[..., 1:] - bi[..., :-1]
        second_dev = torch.linalg.solve(matrix, vi)# [b, n-1]
        second_dev = torch.cat([torch.zeros_like(second_dev[..., :1]), second_dev, torch.zeros_like(second_dev[..., :1])], dim = -1)# [b, n+1]

        coeff = torch.zeros([batch_size, pts_num - 1, 4], device=bins.device)
        coeff[..., 3] = (second_dev[..., 1:] - second_dev[..., :-1]) / (6 * bin_length)
        coeff[..., 2] = second_dev[..., :-1] / 2
        coeff[..., 1] = -bin_length*second_dev[..., 1:]/6 - bin_length*second_dev[..., :-1]/3 + (yy[..., 1:] - yy[..., :-1]) / bin_length
        coeff[..., 0] = yy[..., :-1]
    
        return coeff, degree
//  [python]
def _setup_linear(self):
        """Create parameters that allow fast linear stepping of state"""
        w = self._w()
        B = _r2c(self.B)  # (H N)
        P = _r2c(self.P)
        Q = P.conj()

        # Repeat w shape properly
        B = repeat(B, "1 t n -> 1 (v t) n", v=self.repeat)
        P = repeat(P, "r t n -> r (v t) n", v=self.repeat)
        Q = repeat(Q, "r t n -> r (v t) n", v=self.repeat)
        w = repeat(w, "t n -> (v t) n", v=self.repeat)

        # Prepare Linear stepping
        dt = torch.exp(self.log_dt)
        D = (2.0 / dt.unsqueeze(-1) - w).reciprocal()  # (H, N)
        R = (
            torch.eye(self.rank, dtype=w.dtype, device=w.device)
            + 2 * contract("r h n, h n, s h n -> h r s", Q, D, P).real
        )  # (H R R)
        Q_D = rearrange(Q * D, "r h n -> h r n")
        try:
            R = torch.solve(Q_D, R)[0]  # (H R N)
        except:
            R = torch.tensor(
                np.linalg.solve(
                    R.to(Q_D).contiguous().detach().cpu(),
                    Q_D.contiguous().detach().cpu(),
                )
            ).to(Q_D)
        R = rearrange(R, "h r n -> r h n")

        self.step_params = {
            "D": D,  # (H N)
            "R": R,  # (R H N)
            "P": P,  # (R H N)
            "Q": Q,  # (R H N)
            "B": B,  # (1 H N)
            "E": 2.0 / dt.unsqueeze(-1) + w,  # (H N)
        }
//  [python]
def _setup_linear(self):
        """Create parameters that allow fast linear stepping of state"""
        w = self._w()
        B = _r2c(self.B)  # (H N)
        P = _r2c(self.P)
        Q = P.conj()

        # Repeat w shape properly
        B = repeat(B, "1 t n -> 1 (v t) n", v=self.repeat)
        P = repeat(P, "r t n -> r (v t) n", v=self.repeat)
        Q = repeat(Q, "r t n -> r (v t) n", v=self.repeat)
        w = repeat(w, "t n -> (v t) n", v=self.repeat)

        # Prepare Linear stepping
        dt = torch.exp(self.log_dt)
        D = (2.0 / dt.unsqueeze(-1) - w).reciprocal()  # (H, N)
        R = (
            torch.eye(self.rank, dtype=w.dtype, device=w.device)
            + 2 * contract("r h n, h n, s h n -> h r s", Q, D, P).real
        )  # (H R R)
        Q_D = rearrange(Q * D, "r h n -> h r n")
        try:
            R = torch.linalg.solve(R, Q_D)  # (H R N)
        except:
            R = torch.tensor(
                np.linalg.solve(
                    R.to(Q_D).contiguous().detach().cpu(),
                    Q_D.contiguous().detach().cpu(),
                )
            ).to(Q_D)
        R = rearrange(R, "h r n -> r h n")

        self.step_params = {
            "D": D,  # (H N)
            "R": R,  # (R H N)
            "P": P,  # (R H N)
            "Q": Q,  # (R H N)
            "B": B,  # (1 H N)
            "E": 2.0 / dt.unsqueeze(-1) + w,  # (H N)
        }
//  [python]
def solve_adjoint(self, fwd_sols, param_vals, times):
        #only works for implicit backward euler
        adj_sols = torch.empty_like(fwd_sols)
        ndof, ntimes = fwd_sols.shape
        dqdu = self._dqdu(fwd_sols, ntimes-1, param_vals)
        assert np.allclose(dqdu[:, 0], 0)  # should be zero for implicit euler
        Id = torch.eye(adj_sols.shape[0])
        for ii in range(ntimes-1, 0, -1):
            raw_jac = self.physics._transient_residual(
                fwd_sols[:, ii], times[ii])[1]
            deltat = times[ii]-times[ii-1]
            # rhs = adj_sols[:, ii+1]+deltat*(-dqdu[:, ii+1])
            # print(ii, times[ii])
            # print(dqdu[:, ii])
            if ii < ntimes-1:
                rhs = adj_sols[:, ii+1]-dqdu[:, ii]
            else:
                rhs = -dqdu[:, ii]
            jac_adjoint = self.physics.mesh._apply_boundary_conditions(
                self.physics._bndry_conds, rhs,
                Id-deltat*raw_jac, fwd_sols[:, ii],
                self.physics.flux_jac)[1].T
            # print('jT', jac_adjoint.T)
            adj_sols[:, ii] = torch.solve(rhs, jac_adjoint)[0]
            # print(adj_sols[:, ii], 'phi', adj_sols[:, ii].sum())
        adj_sols[:, 0] = adj_sols[:, 1]
        return adj_sols
//  [python]
def solve_adjoint(self, fwd_sols, param_vals, times):
        #only works for implicit backward euler
        adj_sols = torch.empty_like(fwd_sols)
        ndof, ntimes = fwd_sols.shape
        dqdu = self._dqdu(fwd_sols, ntimes-1, param_vals)
        assert np.allclose(dqdu[:, 0], 0)  # should be zero for implicit euler
        Id = torch.eye(adj_sols.shape[0])
        for ii in range(ntimes-1, 0, -1):
            raw_jac = self.physics._transient_residual(
                fwd_sols[:, ii], times[ii])[1]
            deltat = times[ii]-times[ii-1]
            # rhs = adj_sols[:, ii+1]+deltat*(-dqdu[:, ii+1])
            # print(ii, times[ii])
            # print(dqdu[:, ii])
            if ii < ntimes-1:
                rhs = adj_sols[:, ii+1]-dqdu[:, ii]
            else:
                rhs = -dqdu[:, ii]
            jac_adjoint = self.physics.mesh._apply_boundary_conditions(
                self.physics._bndry_conds, rhs,
                Id-deltat*raw_jac, fwd_sols[:, ii],
                self.physics.flux_jac)[1].T
            # print('jT', jac_adjoint.T)
            adj_sols[:, ii] = torch.linalg.solve(jac_adjoint, rhs)
            # print(adj_sols[:, ii], 'phi', adj_sols[:, ii].sum())
        adj_sols[:, 0] = adj_sols[:, 1]
        return adj_sols
//  [python]
def _solve_least_squares(self):
        solution = torch.solve(self._Q.T @ self._b, self._R)[0]
        self._solution = solution
        self._mu = solution[0]
        self._phi = solution[1:]
//  [python]
def _solve_least_squares(self):
        solution = torch.linalg.solve(self._R, self._Q.T @ self._b)
        self._solution = solution
        self._mu = solution[0]
        self._phi = solution[1:]
//  [python]
def test_caching_parametrization(self):
        r"""Test the caching system of a parametrization"""
        # Define a couple matrix parametrizations
        class Skew(nn.Module):
            def forward(self, X):
                X = X.tril(-1)
                return X - X.T

        class Orthogonal(nn.Module):
            def forward(self, X):
                Id = torch.eye(X.size(0), device=X.device)
                return torch.solve(Id - X, Id + X)[0]

        model = nn.Linear(5, 5)
        parametrize.register_parametrization(model, "weight", Skew())
        parametrize.register_parametrization(model, "weight", Orthogonal())

        # Test that the caching system works
        with parametrize.cached():
            X = model.weight
            Y = model.weight
            self.assertEqual(id(X), id(Y))
//  [python]
def test_caching_parametrization(self):
        r"""Test the caching system of a parametrization"""
        # Define a couple matrix parametrizations
        class Skew(nn.Module):
            def forward(self, X):
                X = X.tril(-1)
                return X - X.T

        class Orthogonal(nn.Module):
            def forward(self, X):
                Id = torch.eye(X.size(0), device=X.device)
                return torch.linalg.solve(Id + X, Id - X)

        model = nn.Linear(5, 5)
        parametrize.register_parametrization(model, "weight", Skew())
        parametrize.register_parametrization(model, "weight", Orthogonal())

        # Test that the caching system works
        with parametrize.cached():
            X = model.weight
            Y = model.weight
            self.assertEqual(id(X), id(Y))