jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def solve_thunks(self,
                     make_a: Callable[[], torch.Tensor],
                     make_b: Callable[[], torch.Tensor]) -> torch.Tensor:
        # We want the least nonnegative solution of (I-a)x = b, and
        # want to use torch.linalg.solve if we can, but there are a
        # number of things that can go wrong:
        # - If a has an eigenvalue = 1, torch.linalg.solve raises RuntimeError.
        # - If a has an eigenvalue > 1, the solution will have negative components.
        # - If a has an eigenvalue = inf, the solution will have -0.0 components.
        # In these cases, we have to fall back to Semiring.solve.
        a = make_a()
        b = make_b()
        try:
            a.neg_().diagonal().add_(1)
            x = torch.solve(b, a)[0]
            if torch.all(torch.copysign(torch.tensor(1.), x) > 0.): # catches -0.0
                return x
        except RuntimeError as e:
            if '(Cannot allocate memory)' in str(e): raise
            pass
        return Semiring.solve_thunks(self, make_a, make_b)
//  [python]
def solve_thunks(self,
                     make_a: Callable[[], torch.Tensor],
                     make_b: Callable[[], torch.Tensor]) -> torch.Tensor:
        # We want the least nonnegative solution of (I-a)x = b, and
        # want to use torch.linalg.solve if we can, but there are a
        # number of things that can go wrong:
        # - If a has an eigenvalue = 1, torch.linalg.solve raises RuntimeError.
        # - If a has an eigenvalue > 1, the solution will have negative components.
        # - If a has an eigenvalue = inf, the solution will have -0.0 components.
        # In these cases, we have to fall back to Semiring.solve.
        a = make_a()
        b = make_b()
        try:
            a.neg_().diagonal().add_(1)
            x = torch.linalg.solve(a, b)
            if torch.all(torch.copysign(torch.tensor(1.), x) > 0.): # catches -0.0
                return x
        except RuntimeError as e:
            if '(Cannot allocate memory)' in str(e): raise
            pass
        return Semiring.solve_thunks(self, make_a, make_b)
//  [python]
def _get_steering_vector_power(
        self, psd_s: torch.Tensor, psd_n: torch.Tensor, reference_vector: torch.Tensor
    ) -> torch.Tensor:
        r"""Estimate the steering vector by the power method.

        Args:
            psd_s (torch.tensor): covariance matrix of speech
                Tensor of dimension `(..., freq, channel, channel)`
            psd_n (torch.Tensor): covariance matrix of noise
                Tensor of dimension `(..., freq, channel, channel)`
            reference_vector (torch.Tensor): one-hot reference channel matrix

        Returns:
            torch.Tensor: the enhanced STFT
                Tensor of dimension `(..., freq, channel, 1)`
        """
        phi = torch.solve(psd_s, psd_n)[0]  # psd_n.inv() @ psd_s
        stv = torch.einsum("...fec,...c->...fe", [phi, reference_vector])
        stv = stv.unsqueeze(-1)
        stv = torch.matmul(phi, stv)
        stv = torch.matmul(psd_s, stv)
        return stv
//  [python]
def _get_steering_vector_power(
        self, psd_s: torch.Tensor, psd_n: torch.Tensor, reference_vector: torch.Tensor
    ) -> torch.Tensor:
        r"""Estimate the steering vector by the power method.

        Args:
            psd_s (torch.tensor): covariance matrix of speech
                Tensor of dimension `(..., freq, channel, channel)`
            psd_n (torch.Tensor): covariance matrix of noise
                Tensor of dimension `(..., freq, channel, channel)`
            reference_vector (torch.Tensor): one-hot reference channel matrix

        Returns:
            torch.Tensor: the enhanced STFT
                Tensor of dimension `(..., freq, channel, 1)`
        """
        phi = torch.linalg.solve(psd_n, psd_s)  # psd_n.inv() @ psd_s
        stv = torch.einsum("...fec,...c->...fe", [phi, reference_vector])
        stv = stv.unsqueeze(-1)
        stv = torch.matmul(phi, stv)
        stv = torch.matmul(psd_s, stv)
        return stv
//  [python]
def forward(self, covmat, rets=None):
        """Perform forward pass.

        Parameters
        ----------
        covmat : torch.Tensor
            Covariance matrix of shape `(n_samples, n_assets, n_assets)`.

        rets : torch.Tensor or None
            If tensor then of shape `(n_samples, n_assets)` representing expected returns. If provided triggers
            computation of maximum share ratio. Else None triggers computation of minimum variance portfolio.

        Returns
        -------
        weights : torch.Tensor
            Of shape (n_samples, n_assets) representing the optimal weights. If `rets` provided, then it represents
            maximum sharpe ratio portfolio (tangency portfolio). Otherwise minimum variance portfolio.
        """
        n_samples, n_assets, _ = covmat.shape
        device = covmat.device
        dtype = covmat.dtype
        ones = torch.ones(n_samples, n_assets, 1)
        if rets is not None:
            expected_returns = rets.view(n_samples, n_assets, 1)
        else:
            expected_returns = ones
        w_unscaled = torch.solve(expected_returns, covmat)[0]
        denominator = torch.matmul(ones.permute(0, 2, 1), w_unscaled)
        w = w_unscaled / denominator
        return w.squeeze(-1)
//  [python]
def forward(self, covmat, rets=None):
        """Perform forward pass.

        Parameters
        ----------
        covmat : torch.Tensor
            Covariance matrix of shape `(n_samples, n_assets, n_assets)`.

        rets : torch.Tensor or None
            If tensor then of shape `(n_samples, n_assets)` representing expected returns. If provided triggers
            computation of maximum share ratio. Else None triggers computation of minimum variance portfolio.

        Returns
        -------
        weights : torch.Tensor
            Of shape (n_samples, n_assets) representing the optimal weights. If `rets` provided, then it represents
            maximum sharpe ratio portfolio (tangency portfolio). Otherwise minimum variance portfolio.
        """
        n_samples, n_assets, _ = covmat.shape
        device = covmat.device
        dtype = covmat.dtype
        ones = torch.ones(n_samples, n_assets, 1)
        if rets is not None:
            expected_returns = rets.view(n_samples, n_assets, 1)
        else:
            expected_returns = ones
        w_unscaled = torch.linalg.solve(covmat, expected_returns)
        denominator = torch.matmul(ones.permute(0, 2, 1), w_unscaled)
        w = w_unscaled / denominator
        return w.squeeze(-1)
//  [python]
def _phi_prep(self, t, J):
        """Construct `V` and its inverse given `t` and couplings `J`."""
        assert t.ndim == 2, f'Tensor `t` should have either shape (batch, 1) or (batch, N) but found shape {t.shape}'
        t = t.repeat(1, self.num_spins) if t.size(-1) == 1 else t
        V = torch.diag_embed(t) - J
        V_inv = torch.solve(batch_eye_like(V), V)[1]
        return t, V, V_inv
//  [python]
def _phi_prep(self, t, J):
        """Construct `V` and its inverse given `t` and couplings `J`."""
        assert t.ndim == 2, f'Tensor `t` should have either shape (batch, 1) or (batch, N) but found shape {t.shape}'
        t = t.repeat(1, self.num_spins) if t.size(-1) == 1 else t
        V = torch.diag_embed(t) - J
        V_inv = torch.linalg.solve(V, batch_eye_like(V))
        return t, V, V_inv
//  [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.broadcast)
        P = repeat(P, "r t n -> r (v t) n", v=self.broadcast)
        Q = repeat(Q, "r t n -> r (v t) n", v=self.broadcast)
        w = repeat(w, "t n -> (v t) n", v=self.broadcast)

        # 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)[1]  # (H R N)
        except:  # noqa
            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.broadcast)
        P = repeat(P, "r t n -> r (v t) n", v=self.broadcast)
        Q = repeat(Q, "r t n -> r (v t) n", v=self.broadcast)
        w = repeat(w, "t n -> (v t) n", v=self.broadcast)

        # 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:  # noqa
            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 fit_causal_prediction(
    X_train, X_test, swave, lam=1e-3, tlag=1, device=torch.device("cuda")
):
    """predict X in the future with exponential filters"""
    # fit on train data
    Xfilt = causal_filter(X_train, swave, tlag, device=device)
    Xfilt = Xfilt.reshape(Xfilt.shape[0], -1)
    NT = X_train.shape[1] * X_train.shape[2]
    nff = Xfilt.shape[0]
    CC = (Xfilt @ Xfilt.T) / NT + lam * torch.eye(nff, device=device)
    CX = (Xfilt @ X_train.reshape(-1, NT).T) / NT
    B = torch.solve(CX, CC)[1]

    # performance on test data
    Xfilt = causal_filter(X_test, swave, tlag, remove_start=True, device=device)
    Xfilt = Xfilt.reshape(Xfilt.shape[0], -1)
    ypred = B.T @ Xfilt
    nt = swave.shape[1]
    ve = compute_varexp(X_test[:, :, nt:].reshape(X_test.shape[0], -1).T, ypred.T)
    return ve, ypred, B
//  [python]
def fit_causal_prediction(
    X_train, X_test, swave, lam=1e-3, tlag=1, device=torch.device("cuda")
):
    """predict X in the future with exponential filters"""
    # fit on train data
    Xfilt = causal_filter(X_train, swave, tlag, device=device)
    Xfilt = Xfilt.reshape(Xfilt.shape[0], -1)
    NT = X_train.shape[1] * X_train.shape[2]
    nff = Xfilt.shape[0]
    CC = (Xfilt @ Xfilt.T) / NT + lam * torch.eye(nff, device=device)
    CX = (Xfilt @ X_train.reshape(-1, NT).T) / NT
    B = torch.linalg.solve(CC, CX)

    # performance on test data
    Xfilt = causal_filter(X_test, swave, tlag, remove_start=True, device=device)
    Xfilt = Xfilt.reshape(Xfilt.shape[0], -1)
    ypred = B.T @ Xfilt
    nt = swave.shape[1]
    ve = compute_varexp(X_test[:, :, nt:].reshape(X_test.shape[0], -1).T, ypred.T)
    return ve, ypred, B
//  [python]
def _get_mvdr_vector(
        self,
        psd_s: torch.Tensor,
        psd_n: torch.Tensor,
        reference_vector: torch.Tensor,
        solution: str = "ref_channel",
        diagonal_loading: bool = True,
        diag_eps: float = 1e-7,
        eps: float = 1e-8,
    ) -> torch.Tensor:
        r"""Compute beamforming vector by the reference channel selection method.

        Args:
            psd_s (torch.Tensor): psd matrix of target speech
            psd_n (torch.Tensor): psd matrix of noise
            reference_vector (torch.Tensor): one-hot reference channel matrix
            solution (str, optional): the solution to estimate the beamforming weight
                (Default: ``ref_channel``)
            diagonal_loading (bool, optional): whether to apply diagonal loading to psd_n
                (Default: ``True``)
            diag_eps (float, optional): The coefficient multipied to the identity matrix for diagonal loading
                (Default: 1e-7)
            eps (float, optional): a value added to the denominator in mask normalization. Default: 1e-8

        Returns:
            torch.Tensor: the mvdr beamforming weight matrix
        """
        if diagonal_loading:
            psd_n = self._tik_reg(psd_n, reg=diag_eps, eps=eps)
        if solution == "ref_channel":
            numerator = torch.solve(psd_s, psd_n)[0]  # psd_n.inv() @ psd_s
            # ws: (..., C, C) / (...,) -> (..., C, C)
            ws = numerator / (_get_mat_trace(numerator)[..., None, None] + eps)
            # h: (..., F, C_1, C_2) x (..., C_2) -> (..., F, C_1)
            beamform_vector = torch.einsum("...fec,...c->...fe", [ws, reference_vector])
        else:
            if solution == "stv_evd":
                stv = self._get_steering_vector_evd(psd_s)
            else:
                stv = self._get_steering_vector_power(psd_s, psd_n, reference_vector)
            # numerator = psd_n.inv() @ stv
            numerator = torch.solve(stv, psd_n)[0].squeeze(-1)  # (..., freq, channel)
            # denominator = stv^H @ psd_n.inv() @ stv
            denominator = torch.einsum("...d,...d->...", [stv.conj().squeeze(-1), numerator])
            # normalzie the numerator
            scale = stv.squeeze(-1)[..., self.ref_channel, None].conj()
            beamform_vector = numerator / (denominator.real.unsqueeze(-1) + eps) * scale

        return beamform_vector
//  [python]
def _get_mvdr_vector(
        self,
        psd_s: torch.Tensor,
        psd_n: torch.Tensor,
        reference_vector: torch.Tensor,
        solution: str = "ref_channel",
        diagonal_loading: bool = True,
        diag_eps: float = 1e-7,
        eps: float = 1e-8,
    ) -> torch.Tensor:
        r"""Compute beamforming vector by the reference channel selection method.

        Args:
            psd_s (torch.Tensor): psd matrix of target speech
            psd_n (torch.Tensor): psd matrix of noise
            reference_vector (torch.Tensor): one-hot reference channel matrix
            solution (str, optional): the solution to estimate the beamforming weight
                (Default: ``ref_channel``)
            diagonal_loading (bool, optional): whether to apply diagonal loading to psd_n
                (Default: ``True``)
            diag_eps (float, optional): The coefficient multipied to the identity matrix for diagonal loading
                (Default: 1e-7)
            eps (float, optional): a value added to the denominator in mask normalization. Default: 1e-8

        Returns:
            torch.Tensor: the mvdr beamforming weight matrix
        """
        if diagonal_loading:
            psd_n = self._tik_reg(psd_n, reg=diag_eps, eps=eps)
        if solution == "ref_channel":
            numerator = torch.linalg.solve(psd_n, psd_s)  # psd_n.inv() @ psd_s
            # ws: (..., C, C) / (...,) -> (..., C, C)
            ws = numerator / (_get_mat_trace(numerator)[..., None, None] + eps)
            # h: (..., F, C_1, C_2) x (..., C_2) -> (..., F, C_1)
            beamform_vector = torch.einsum("...fec,...c->...fe", [ws, reference_vector])
        else:
            if solution == "stv_evd":
                stv = self._get_steering_vector_evd(psd_s)
            else:
                stv = self._get_steering_vector_power(psd_s, psd_n, reference_vector)
            # numerator = psd_n.inv() @ stv
            numerator = torch.linalg.solve(psd_n, stv).squeeze(-1)  # (..., freq, channel)
            # denominator = stv^H @ psd_n.inv() @ stv
            denominator = torch.einsum("...d,...d->...", [stv.conj().squeeze(-1), numerator])
            # normalzie the numerator
            scale = stv.squeeze(-1)[..., self.ref_channel, None].conj()
            beamform_vector = numerator / (denominator.real.unsqueeze(-1) + eps) * scale

        return beamform_vector
//  [python]
def transform(self, results: dict) -> dict:
        # TODO: transforms should use numpy instead of torch
        points = results['points']
        global_points = []
        points_slice_indices = [0]
        for idx in range(len(points)):
            point = points[idx].tensor[..., :3]
            point = torch.cat([point, point.new_ones(point.shape[0], 1)],
                              dim=1)
            global2ego = torch.from_numpy(
                results['depth2img']['extrinsic'][idx]).to(point.device)
            global_point = torch.solve(point.transpose(
                0, 1), global2ego).solution.transpose(0, 1)
            points[idx].tensor[:, :3] = global_point[:, :3]
            global_points.append(points[idx].tensor)
            if self.save_slices:
                points_slice_indices.append(points_slice_indices[-1] +
                                            len(points[idx].tensor))
        points = torch.cat(global_points)
        # a little hard code, to be improved
        points_class = get_points_type(self.coord_type)
        points = points_class(
            points,
            points_dim=results['points'][0].points_dim,
            attribute_dims=results['points'][0].attribute_dims)
        results['points'] = points

        if self.save_slices:
            results['points_slice_indices'] = points_slice_indices

        return results
//  [python]
def transform(self, results: dict) -> dict:
        # TODO: transforms should use numpy instead of torch
        points = results['points']
        global_points = []
        points_slice_indices = [0]
        for idx in range(len(points)):
            point = points[idx].tensor[..., :3]
            point = torch.cat([point, point.new_ones(point.shape[0], 1)],
                              dim=1)
            global2ego = torch.from_numpy(
                results['depth2img']['extrinsic'][idx]).to(point.device)
            global_point = (torch.linalg.solve(global2ego, point.transpose(
                0, 1))).transpose(0, 1)
            points[idx].tensor[:, :3] = global_point[:, :3]
            global_points.append(points[idx].tensor)
            if self.save_slices:
                points_slice_indices.append(points_slice_indices[-1] +
                                            len(points[idx].tensor))
        points = torch.cat(global_points)
        # a little hard code, to be improved
        points_class = get_points_type(self.coord_type)
        points = points_class(
            points,
            points_dim=results['points'][0].points_dim,
            attribute_dims=results['points'][0].attribute_dims)
        results['points'] = points

        if self.save_slices:
            results['points_slice_indices'] = points_slice_indices

        return results
//  [python]
def solve(a, b):
    return torch.solve(b, a)[0]
//  [python]
def solve(a, b):
    return torch.linalg.solve(a, b)
//  [python]
def solve(A: torch.tensor, b: torch.Tensor, type: str):
    # Use pytorch to solve the linear system
    if type == "torch":
        return torch.solve(b, A)[1]

    # Use colrow to solve the linear system
    nrows = A.shape[-1]
    ncols = b.shape[-1]
    dim = cal_dim(A, b)
    resid_step = 10
    max_niter = int(1.5 * nrows)
    rtol = 1e-6
    atol = 1e-6
    eps = 1e-8

    if torch.allclose(b, b * 0, rtol=rtol, atol=atol):
        x0 = torch.zeros((*dim, nrows, ncols), dtype=A.dtype, device=A.device)
        return x0

    presec_func = lambda x: x
    hermit_base = True
    weight_trans_func, value_trans, col_swap = decompose(A, b, dim, hermit_base)
    value_norm = value_trans.norm(dim=-2, keepdim=True)
    stop_matrix = torch.max(rtol * value_norm, atol * torch.ones_like(value_norm))

    weight_shape = (ncols, *dim, nrows, 1) if col_swap else (*dim, nrows, ncols)

    x_sim = torch.zeros(weight_shape, dtype=A.dtype, device=A.device)
    r_sim = value_trans - weight_trans_func(x_sim)
    z_sim = presec_func(r_sim)
    p_sim = z_sim
    risk = transf(r_sim, z_sim)
    best_resid = r_sim.norm(dim=-2).max().item()
    best_x_sim = x_sim

    for k in range(1, max_niter + 1):
        weight_sim = weight_trans_func(p_sim)
        alphak = risk / norm(transf(p_sim, weight_sim), eps)
        x_sim2 = x_sim + alphak * p_sim

        if resid_step != 0 and k % resid_step == 0:
            r_sim2 = value_trans - weight_trans_func(x_sim2)
        else:
            r_sim2 = r_sim - alphak * weight_sim

        resid = r_sim2
        resid_norm = resid.norm(dim=-2, keepdim=True)

        max_resid_norm = resid_norm.max().item()
        if max_resid_norm < best_resid:
            best_resid = max_resid_norm
            best_x_sim = x_sim2

        z_sim2 = presec_func(r_sim2)
        risk2 = transf(r_sim2, z_sim2)
        b_risk = risk2 / norm(risk, eps)
        p_sim2 = z_sim2 + b_risk * p_sim

        p_sim = p_sim2
        z_sim = z_sim2
        x_sim = x_sim2
        r_sim = r_sim2
        risk = risk2

    x_sim2 = best_x_sim
    if col_swap:
        x_sim2 = x_sim2.transpose(0, -1).squeeze(0)
    return x_sim2
//  [python]
def solve(A: torch.tensor, b: torch.Tensor, type: str):
    # Use pytorch to solve the linear system
    if type == "torch":
        return torch.linalg.solve(A, b)

    # Use colrow to solve the linear system
    nrows = A.shape[-1]
    ncols = b.shape[-1]
    dim = cal_dim(A, b)
    resid_step = 10
    max_niter = int(1.5 * nrows)
    rtol = 1e-6
    atol = 1e-6
    eps = 1e-8

    if torch.allclose(b, b * 0, rtol=rtol, atol=atol):
        x0 = torch.zeros((*dim, nrows, ncols), dtype=A.dtype, device=A.device)
        return x0

    presec_func = lambda x: x
    hermit_base = True
    weight_trans_func, value_trans, col_swap = decompose(A, b, dim, hermit_base)
    value_norm = value_trans.norm(dim=-2, keepdim=True)
    stop_matrix = torch.max(rtol * value_norm, atol * torch.ones_like(value_norm))

    weight_shape = (ncols, *dim, nrows, 1) if col_swap else (*dim, nrows, ncols)

    x_sim = torch.zeros(weight_shape, dtype=A.dtype, device=A.device)
    r_sim = value_trans - weight_trans_func(x_sim)
    z_sim = presec_func(r_sim)
    p_sim = z_sim
    risk = transf(r_sim, z_sim)
    best_resid = r_sim.norm(dim=-2).max().item()
    best_x_sim = x_sim

    for k in range(1, max_niter + 1):
        weight_sim = weight_trans_func(p_sim)
        alphak = risk / norm(transf(p_sim, weight_sim), eps)
        x_sim2 = x_sim + alphak * p_sim

        if resid_step != 0 and k % resid_step == 0:
            r_sim2 = value_trans - weight_trans_func(x_sim2)
        else:
            r_sim2 = r_sim - alphak * weight_sim

        resid = r_sim2
        resid_norm = resid.norm(dim=-2, keepdim=True)

        max_resid_norm = resid_norm.max().item()
        if max_resid_norm < best_resid:
            best_resid = max_resid_norm
            best_x_sim = x_sim2

        z_sim2 = presec_func(r_sim2)
        risk2 = transf(r_sim2, z_sim2)
        b_risk = risk2 / norm(risk, eps)
        p_sim2 = z_sim2 + b_risk * p_sim

        p_sim = p_sim2
        z_sim = z_sim2
        x_sim = x_sim2
        r_sim = r_sim2
        risk = risk2

    x_sim2 = best_x_sim
    if col_swap:
        x_sim2 = x_sim2.transpose(0, -1).squeeze(0)
    return x_sim2
//  [python]
def torch_solve_qp_uncon(Q, p):
    #######################################################################
    # Solve a QP in form:
    #   x_star =   argmin_x 0.5*x^TQx + p^Tx

    # Q:  A (n_batch,n_x,n_x) SPD matrix
    # p:  A (n_batch,n_x,1) matrix.
    # Returns: x_star:  A (n_x,1) matrix
    #######################################################################
    x = torch.solve(-p, Q)[0]
    sol = {'x': x}
    return sol
//  [python]
def torch_solve_qp_uncon(Q, p):
    #######################################################################
    # Solve a QP in form:
    #   x_star =   argmin_x 0.5*x^TQx + p^Tx

    # Q:  A (n_batch,n_x,n_x) SPD matrix
    # p:  A (n_batch,n_x,1) matrix.
    # Returns: x_star:  A (n_x,1) matrix
    #######################################################################
    x = torch.linalg.solve(Q, -p)
    sol = {'x': x}
    return sol
//  [python]
def update_u(self, x: torch.Tensor, u: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
        # x ≈ u @ t(v) --> u = ?
        *_, M, N = x.shape
        if M >= N:
            u_new = x @ t(torch.linalg.pinv(v))
        else:
            a, b = x @ v, t(v) @ v
            u_new = t(torch.solve(t(a), b)[0])

        return self.project(u_new)
//  [python]
def update_u(self, x: torch.Tensor, u: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
        # x ≈ u @ t(v) --> u = ?
        *_, M, N = x.shape
        if M >= N:
            u_new = x @ t(torch.linalg.pinv(v))
        else:
            a, b = x @ v, t(v) @ v
            u_new = torch.linalg.solve(b, t(a))
            u_new = t(u_new)

        return self.project(u_new)
//  [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 cayley_map(X):
    n = X.size(-1)
    Id = torch.eye(n, dtype=X.dtype, device=X.device)
    return torch.solve(Id.add(X, alpha=0.5), Id.add(X, alpha=-0.5))[0]
//  [python]
def cayley_map(X):
    n = X.size(-1)
    Id = torch.eye(n, dtype=X.dtype, device=X.device)
    return torch.linalg.solve(Id.add(X, alpha=-0.5), Id.add(X, alpha=0.5))
//  [python]
def solveAlphas_(pw,control_w):
    b = pw.shape[0]
    n = pw.shape[1]

    ones = torch.ones((b,n,1)).to(pw.device)
    ph_w = torch.cat([pw,ones],dim=-1)

    alphas = torch.solve(ph_w.permute(0,2,1),control_w)[1]

    return alphas
//  [python]
def solveAlphas_(pw,control_w):
    b = pw.shape[0]
    n = pw.shape[1]

    ones = torch.ones((b,n,1)).to(pw.device)
    ph_w = torch.cat([pw,ones],dim=-1)

    alphas = torch.linalg.solve(control_w,ph_w.permute(0,2,1))

    return alphas
//  [python]
def __init__(self, sf, ell, X, y, eps=1e-5, kernel='exp'):
        self.sf = sf
        self.ell = ell
        self.X = X
        self.y = y
        self.eps = eps
        self.K = K if kernel=='exp' else Klinear
        self.KXX_inv_y = torch.solve(self.K(X,X,ell,sf,eps), y)[0]
//  [python]
def __init__(self, sf, ell, X, y, eps=1e-5, kernel='exp'):
        self.sf = sf
        self.ell = ell
        self.X = X
        self.y = y
        self.eps = eps
        self.K = K if kernel=='exp' else Klinear
        self.KXX_inv_y = torch.linalg.solve(self.K(X,X,ell,sf,eps), y)[0]
//  [python]
def likelihood(self, G, batch):
        x, dx, gc = batch
        # Reshape x, dx like they have an additional time dimension
        # [batch, time, dim]
        Gt = torch.transpose(G.to(x), -2, -1).unsqueeze(1)
        x = x.unsqueeze(1).unsqueeze(0)
        dx = dx.unsqueeze(1)
        x_masked = Gt * x
        A_est = []
        for p in range(self.n_dim):
            w_est = torch.solve(
                torch.transpose(x_masked[:, :, p, :], -2, -1) @ dx[:, :, p],
                (torch.transpose(x_masked[:, :, p, :], -2, -1) @ x_masked[:, :, p, :])
                + self.beta * torch.eye(self.n_dim).unsqueeze(0).type_as(x_masked),
            )[0]
            A_est.append(w_est)
        A_est = torch.cat(A_est, dim=2)
        A_est = A_est.unsqueeze(1)
        pred = A_est @ torch.transpose(x, -2, -1)
        pred = torch.transpose(pred, -2, -1)
        mses = nn.functional.mse_loss(pred, dx.expand(*pred.shape), reduction="none")
        mses = mses.mean((1, 2, 3))
        return mses
//  [python]
def likelihood(self, G, batch):
        x, dx, gc = batch
        # Reshape x, dx like they have an additional time dimension
        # [batch, time, dim]
        Gt = torch.transpose(G.to(x), -2, -1).unsqueeze(1)
        x = x.unsqueeze(1).unsqueeze(0)
        dx = dx.unsqueeze(1)
        x_masked = Gt * x
        A_est = []
        for p in range(self.n_dim):
            w_est = torch.linalg.solve(
                (torch.transpose(x_masked[:, :, p, :], -2, -1) @ x_masked[:, :, p, :])
                + self.beta * torch.eye(self.n_dim).unsqueeze(0).type_as(x_masked),
                torch.transpose(x_masked[:, :, p, :], -2, -1) @ dx[:, :, p],
            )
            A_est.append(w_est)
        A_est = torch.cat(A_est, dim=2)
        A_est = A_est.unsqueeze(1)
        pred = A_est @ torch.transpose(x, -2, -1)
        pred = torch.transpose(pred, -2, -1)
        mses = nn.functional.mse_loss(pred, dx.expand(*pred.shape), reduction="none")
        mses = mses.mean((1, 2, 3))
        return mses
//  [python]
def smw_inv(x, damping=1e-7):
    n, d = x.shape  # n x d
    I = torch.eye(d, device=x.device)
    G = x @ x.T  # n x n (Gram matrix)
    diag = torch.diagonal(G)
    diag += damping * n
    Ginv_x = torch.solve(x, G)[0]  # n x d
    xt_Ginv_x = x.T @ Ginv_x  # d x d
    return (I - xt_Ginv_x) / damping
//  [python]
def smw_inv(x, damping=1e-7):
    n, d = x.shape  # n x d
    I = torch.eye(d, device=x.device)
    G = x @ x.T  # n x n (Gram matrix)
    diag = torch.diagonal(G)
    diag += damping * n
    Ginv_x = torch.linalg.solve(G, x)  # n x d
    xt_Ginv_x = x.T @ Ginv_x  # d x d
    return (I - xt_Ginv_x) / damping
//  [python]
def GP_noise(y1, K11, K12, K22, epsilon_noise, device=torch.device("cpu")):
    """
    Source: https://github.com/PrincetonLIPS/AHGP Licensed under the MIT License
    Calculate the posterior mean and covariance matrix for y2 based on the noisy observations y1 and the given kernel matrix
    """
    # y1: N_train x 1
    # K11: N_train x N_train
    # K12: N_train x N_test
    # K22: N_test x N_test
    # Kernel of the noisy observations
    K11 = K11 + epsilon_noise * torch.eye(K11.shape[0]).to(device)
    solved = torch.solve(K12, K11)[1]
    # Compute posterior mean
    mu_2 = torch.matmul(solved.T, y1)
    var_2 = K22 - torch.matmul(solved.T, K12)
    return mu_2, var_2
//  [python]
def GP_noise(y1, K11, K12, K22, epsilon_noise, device=torch.device("cpu")):
    """
    Source: https://github.com/PrincetonLIPS/AHGP Licensed under the MIT License
    Calculate the posterior mean and covariance matrix for y2 based on the noisy observations y1 and the given kernel matrix
    """
    # y1: N_train x 1
    # K11: N_train x N_train
    # K12: N_train x N_test
    # K22: N_test x N_test
    # Kernel of the noisy observations
    K11 = K11 + epsilon_noise * torch.eye(K11.shape[0]).to(device)
    solved = torch.linalg.solve(K11, K12)
    # Compute posterior mean
    mu_2 = torch.matmul(solved.T, y1)
    var_2 = K22 - torch.matmul(solved.T, K12)
    return mu_2, var_2
//  [python]
def get_geometry(
        self,
        fH,
        fW,
        lidar2img,
        img_metas,
    ):
        B, N, _, _ = lidar2img.shape
        device = lidar2img.device
        if self.frustum == None:
            self.frustum = self.create_frustum(fH,fW,img_metas)
            self.frustum = self.frustum.to(device)
            # self.D = self.frustum.shape[0]
        
        points = self.frustum.view(1,1,self.D, fH, fW, 3) \
                 .repeat(B,N,1,1,1,1)
        lidar2img = lidar2img.view(B,N,1,1,1,4,4)
        # img2lidar = torch.inverse(lidar2img)
        points = torch.cat(
            (points, torch.ones_like(points[..., :1])), -1)
        points = torch.solve(points.unsqueeze(-1).to(torch.float32), 
                                    lidar2img.to(torch.float32)).solution.squeeze(-1)
        # points = torch.matmul(img2lidar.to(torch.float32),
        #                       points.unsqueeze(-1).to(torch.float32)).squeeze(-1)
        eps = 1e-5
        points = points[..., 0:3] / torch.maximum(
            points[..., 3:4], torch.ones_like(points[..., 3:4]) * eps)

        return points
//  [python]
def get_geometry(
        self,
        fH,
        fW,
        lidar2img,
        img_metas,
    ):
        B, N, _, _ = lidar2img.shape
        device = lidar2img.device
        if self.frustum == None:
            self.frustum = self.create_frustum(fH,fW,img_metas)
            self.frustum = self.frustum.to(device)
            # self.D = self.frustum.shape[0]
        
        points = self.frustum.view(1,1,self.D, fH, fW, 3) \
                 .repeat(B,N,1,1,1,1)
        lidar2img = lidar2img.view(B,N,1,1,1,4,4)
        # img2lidar = torch.inverse(lidar2img)
        points = torch.cat(
            (points, torch.ones_like(points[..., :1])), -1)
        points = torch.linalg.solve(lidar2img.to(torch.float32), 
                                    points.unsqueeze(-1).to(torch.float32)).squeeze(-1)
        # points = torch.matmul(img2lidar.to(torch.float32),
        #                       points.unsqueeze(-1).to(torch.float32)).squeeze(-1)
        eps = 1e-5
        points = points[..., 0:3] / torch.maximum(
            points[..., 3:4], torch.ones_like(points[..., 3:4]) * eps)

        return points
//  [python]
def _find_coeffs(orig_pts:Points, targ_pts:Points)->Tensor:
    "Find 8 coeff mentioned [here](https://web.archive.org/web/20150222120106/xenia.media.mit.edu/~cwren/interpolator/)."
    matrix = []
    #The equations we'll need to solve.
    for p1, p2 in zip(targ_pts, orig_pts):
        matrix.append([p1[0], p1[1], 1, 0, 0, 0, -p2[0]*p1[0], -p2[0]*p1[1]])
        matrix.append([0, 0, 0, p1[0], p1[1], 1, -p2[1]*p1[0], -p2[1]*p1[1]])

    A = FloatTensor(matrix)
    B = FloatTensor(orig_pts).view(8, 1)
    #The 8 scalars we seek are solution of AX = B
    return torch.solve(B, A)[0][:,0]
//  [python]
def _find_coeffs(orig_pts:Points, targ_pts:Points)->Tensor:
    "Find 8 coeff mentioned [here](https://web.archive.org/web/20150222120106/xenia.media.mit.edu/~cwren/interpolator/)."
    matrix = []
    #The equations we'll need to solve.
    for p1, p2 in zip(targ_pts, orig_pts):
        matrix.append([p1[0], p1[1], 1, 0, 0, 0, -p2[0]*p1[0], -p2[0]*p1[1]])
        matrix.append([0, 0, 0, p1[0], p1[1], 1, -p2[1]*p1[0], -p2[1]*p1[1]])

    A = FloatTensor(matrix)
    B = FloatTensor(orig_pts).view(8, 1)
    #The 8 scalars we seek are solution of AX = B
    return torch.linalg.solve(A,B)[:,0]
//  [python]
def likelihood(self, G, batch, node_idx=None):
        x, dx, gc = batch
        # Reshape x, dx like they have an additional time dimension
        # [batch, time, dim]
        Gt = torch.transpose(G.to(x), -2, -1).unsqueeze(1)
        x = x.unsqueeze(1).unsqueeze(0)
        dx = dx.unsqueeze(1)
        if G.shape[1] == 1:
            x_masked = G.unsqueeze(1) * x[:, :, :, node_idx].unsqueeze(-1)
            A_est = []
            w_est = torch.solve(
                torch.transpose(x_masked[:, :, 0, :], -2, -1) @ dx[:, 0, node_idx],
                (torch.transpose(x_masked[:, :, 0, :], -2, -1) @ x_masked[:, :, 0, :])
                + self.beta * torch.eye(self.n_dim).unsqueeze(0).type_as(x_masked),
            )[0]
            A_est.append(w_est)
            A_est = torch.stack(A_est, dim=2)
            A_est = A_est.unsqueeze(1)
            A_est = torch.transpose(A_est, -1, -2)
            pred = A_est @ torch.transpose(x, -2, -1)
            pred = torch.transpose(pred, -2, -1)
            mses = nn.functional.mse_loss(
                pred,
                dx[:, :, node_idx].unsqueeze(-1).expand(*pred.shape),
                reduction="none",
            )
            mses = mses.mean((1, 2, 3))
        else:
            x_masked = Gt * x
            A_est = []
            for p in range(self.n_dim):
                w_est = torch.solve(
                    torch.transpose(x_masked[:, :, p, :], -2, -1) @ dx[:, :, p],
                    (
                        torch.transpose(x_masked[:, :, p, :], -2, -1)
                        @ x_masked[:, :, p, :]
                    )
                    + self.beta * torch.eye(self.n_dim).unsqueeze(0).type_as(x_masked),
                )[0]
                A_est.append(w_est)
            A_est = torch.cat(A_est, dim=2)
            A_est = A_est.unsqueeze(1)
            pred = A_est @ torch.transpose(x, -2, -1)
            pred = torch.transpose(pred, -2, -1)
            mses = nn.functional.mse_loss(
                pred, dx.expand(*pred.shape), reduction="none"
            )
            mses = mses.mean((1, 2, 3))
        return mses
//  [python]
def likelihood(self, G, batch, node_idx=None):
        x, dx, gc = batch
        # Reshape x, dx like they have an additional time dimension
        # [batch, time, dim]
        Gt = torch.transpose(G.to(x), -2, -1).unsqueeze(1)
        x = x.unsqueeze(1).unsqueeze(0)
        dx = dx.unsqueeze(1)
        if G.shape[1] == 1:
            x_masked = G.unsqueeze(1) * x[:, :, :, node_idx].unsqueeze(-1)
            A_est = []
            w_est = torch.linalg.solve(
                (torch.transpose(x_masked[:, :, 0, :], -2, -1) @ x_masked[:, :, 0, :])
                + self.beta * torch.eye(self.n_dim).unsqueeze(0).type_as(x_masked),
                torch.transpose(x_masked[:, :, 0, :], -2, -1) @ dx[:, 0, node_idx],
            )
            A_est.append(w_est)
            A_est = torch.stack(A_est, dim=2)
            A_est = A_est.unsqueeze(1)
            A_est = torch.transpose(A_est, -1, -2)
            pred = A_est @ torch.transpose(x, -2, -1)
            pred = torch.transpose(pred, -2, -1)
            mses = nn.functional.mse_loss(
                pred,
                dx[:, :, node_idx].unsqueeze(-1).expand(*pred.shape),
                reduction="none",
            )
            mses = mses.mean((1, 2, 3))
        else:
            x_masked = Gt * x
            A_est = []
            for p in range(self.n_dim):
                w_est = torch.linalg.solve(
                    (
                        torch.transpose(x_masked[:, :, p, :], -2, -1)
                        @ x_masked[:, :, p, :]
                    )
                    + self.beta * torch.eye(self.n_dim).unsqueeze(0).type_as(x_masked),
                    torch.transpose(x_masked[:, :, p, :], -2, -1) @ dx[:, :, p],
                )
                A_est.append(w_est)
            A_est = torch.cat(A_est, dim=2)
            A_est = A_est.unsqueeze(1)
            pred = A_est @ torch.transpose(x, -2, -1)
            pred = torch.transpose(pred, -2, -1)
            mses = nn.functional.mse_loss(
                pred, dx.expand(*pred.shape), reduction="none"
            )
            mses = mses.mean((1, 2, 3))
        return mses
//  [python]
def forward(self, x):
        """compute prediction. fit() must have been called.
        x: test input data point. N x D tensor for the data dimensionality D."""
        L = self.L
        sqrt_W = self.sqrt_W
        k = self.kernel_mat(self.X, x)
        v = torch.solve(sqrt_W.mm(k), L)[0]
        mu = k.T.mm(self.likelihood_func.get_jacobian_of_log_likelihood(self.y, self.f))  # (N',1)
        var = self.amplitude_scale - torch.diag(v.T.mm(v))  # (N')

        z = mu.repeat(1, self.n_samples) + torch.sqrt(var).unsqueeze(-1) * torch.randn_like(
            mu.repeat(1, self.n_samples)
        )  # (N',{self.n_samples})

        pi = torch.sigmoid(z).mean(-1)
        return mu, var, pi
//  [python]
def forward(self, x):
        """compute prediction. fit() must have been called.
        x: test input data point. N x D tensor for the data dimensionality D."""
        L = self.L
        sqrt_W = self.sqrt_W
        k = self.kernel_mat(self.X, x)
        v = torch.linalg.solve(L, sqrt_W.mm(k))
        mu = k.T.mm(self.likelihood_func.get_jacobian_of_log_likelihood(self.y, self.f))  # (N',1)
        var = self.amplitude_scale - torch.diag(v.T.mm(v))  # (N')

        z = mu.repeat(1, self.n_samples) + torch.sqrt(var).unsqueeze(-1) * torch.randn_like(
            mu.repeat(1, self.n_samples)
        )  # (N',{self.n_samples})

        pi = torch.sigmoid(z).mean(-1)
        return mu, var, pi
//  [python]
def translation_with(self, tgt_pos):
        """Compute the translation with an other pos

        Parameters
        ----------
        tgt_pos: aloscene.Pose
            Target pose to compute the translation with

        Returns
        -------
        n_pos: torch.tensor
            Translation tensor of shape (..., 3)
        """
        Ttgt2self = torch.solve(tgt_pos.as_tensor(), self.as_tensor())[0]
        translation = Ttgt2self[..., :3, -1]
        return translation
//  [python]
def translation_with(self, tgt_pos):
        """Compute the translation with an other pos

        Parameters
        ----------
        tgt_pos: aloscene.Pose
            Target pose to compute the translation with

        Returns
        -------
        n_pos: torch.tensor
            Translation tensor of shape (..., 3)
        """
        Ttgt2self = torch.linalg.solve(self.as_tensor(), tgt_pos.as_tensor())
        translation = Ttgt2self[..., :3, -1]
        return translation
//  [python]
def Rome(self, k, v, lam=1.5*10**4):
        '''
        get the change of the weight
        :param k: the matrix of the new embedding (num, emb_dim)
        :param v: the matrix of the old embedding (num, emb_dim) and the target for the new embedding (num, target_dim)
        :param old_weight: the old weight of the Linear layer (target_dim, emb_dim)
        :return: the change of the weight (emb_dim, target_dim)
        '''
        old_emb = v
        new_emb = k
        target = v
        old_weight = torch.eye(v.size()[-1]).to(old_emb.device) #Todo add the last layer condition
        C0 = lam * torch.mean(old_emb.unsqueeze(-1) @ old_emb.unsqueeze(-2), dim=0)
        R = target - new_emb @ old_weight.T
        I = torch.eye(C0.shape[0]).to(old_emb.device)
        delta = R.T @ new_emb @ torch.solve(I, C0 + new_emb.T @ new_emb)[0]
        self._open_ROME = True
        self.delta = delta.T
        self.manual_swatch = False
//  [python]
def Rome(self, k, v, lam=1.5*10**4):
        '''
        get the change of the weight
        :param k: the matrix of the new embedding (num, emb_dim)
        :param v: the matrix of the old embedding (num, emb_dim) and the target for the new embedding (num, target_dim)
        :param old_weight: the old weight of the Linear layer (target_dim, emb_dim)
        :return: the change of the weight (emb_dim, target_dim)
        '''
        old_emb = v
        new_emb = k
        target = v
        old_weight = torch.eye(v.size()[-1]).to(old_emb.device) #Todo add the last layer condition
        C0 = lam * torch.mean(old_emb.unsqueeze(-1) @ old_emb.unsqueeze(-2), dim=0)
        R = target - new_emb @ old_weight.T
        I = torch.eye(C0.shape[0]).to(old_emb.device)
        delta = R.T @ new_emb @ torch.linalg.solve(C0 + new_emb.T @ new_emb, I)
        self._open_ROME = True
        self.delta = delta.T
        self.manual_swatch = False
//  [python]
def fit_causal_prediction(X_train, X_test, swave, lam = 1e-3, tlag=1, device=torch.device('cuda')):
    """ predict X in the future with exponential filters"""
    # fit on train data
    Xfilt = causal_filter(X_train, swave, tlag)    
    Xfilt = Xfilt.reshape(Xfilt.shape[0], -1)
    NT = X_train.shape[1] * X_train.shape[2]
    nff = Xfilt.shape[0]
    CC = (Xfilt @ Xfilt.T)/NT + lam * torch.eye(nff, device = device)
    CX = (Xfilt @ X_train.reshape(-1,NT).T) / NT
    B = torch.solve(CX, CC)[0]    

    # performance on test data
    Xfilt = causal_filter(X_test, swave, tlag, remove_start=True)    
    Xfilt = Xfilt.reshape(Xfilt.shape[0], -1)
    ypred = B.T @ Xfilt
    nt = swave.shape[1]
    ve = compute_varexp(X_test[:,:,nt:].reshape(X_test.shape[0],-1).T, ypred.T)
    return ve, ypred, B
//  [python]
def fit_causal_prediction(X_train, X_test, swave, lam = 1e-3, tlag=1, device=torch.device('cuda')):
    """ predict X in the future with exponential filters"""
    # fit on train data
    Xfilt = causal_filter(X_train, swave, tlag)    
    Xfilt = Xfilt.reshape(Xfilt.shape[0], -1)
    NT = X_train.shape[1] * X_train.shape[2]
    nff = Xfilt.shape[0]
    CC = (Xfilt @ Xfilt.T)/NT + lam * torch.eye(nff, device = device)
    CX = (Xfilt @ X_train.reshape(-1,NT).T) / NT
    B = torch.linalg.solve(CC, CX)    

    # performance on test data
    Xfilt = causal_filter(X_test, swave, tlag, remove_start=True)    
    Xfilt = Xfilt.reshape(Xfilt.shape[0], -1)
    ypred = B.T @ Xfilt
    nt = swave.shape[1]
    ve = compute_varexp(X_test[:,:,nt:].reshape(X_test.shape[0],-1).T, ypred.T)
    return ve, ypred, B
//  [python]
def Kinetic(self, X):
      """
      Kinetic energy in the Hamiltonian function

      Parameters
      ----------
      X: torch.Tensor
          training trajectory points in input, with shape [batch size, nop*2s]

      Returns
      -------
      row : torch.Tensor
            Kinetic energy, with shape [batch size, 1]

      """ 

      nop = int(X.size(dim=1)/6)
      id = torch.eye(nop,nop)
      ref = torch.ones(3,3)
      R = torch.kron(id, ref).to(device)
      U = torch.ones(3*nop,3*nop).to(device)
      p = X[:,3:6]
      for j in range(1,nop):
        p = torch.cat([p, X[:,3*(2*j+1):3*(2*j+2)]], axis=1)
      MM = self.MassMat(X) 
      k = (0.5 * torch.einsum('ij,ij->i',p,torch.solve(p, MM)[0])).unsqueeze(1) # q dependent component
      return k
//  [python]
def Kinetic(self, X):
      """
      Kinetic energy in the Hamiltonian function

      Parameters
      ----------
      X: torch.Tensor
          training trajectory points in input, with shape [batch size, nop*2s]

      Returns
      -------
      row : torch.Tensor
            Kinetic energy, with shape [batch size, 1]

      """ 

      nop = int(X.size(dim=1)/6)
      id = torch.eye(nop,nop)
      ref = torch.ones(3,3)
      R = torch.kron(id, ref).to(device)
      U = torch.ones(3*nop,3*nop).to(device)
      p = X[:,3:6]
      for j in range(1,nop):
        p = torch.cat([p, X[:,3*(2*j+1):3*(2*j+2)]], axis=1)
      MM = self.MassMat(X) 
      k = (0.5 * torch.einsum('ij,ij->i',p,torch.linalg.solve(MM,p))).unsqueeze(1) # q dependent component
      return k