jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def fpcm_MOVE_sl(a, env, ctm_args=cfg.ctm_args, global_args=cfg.global_args,
    past_steps_data=None):
    verbosity= ctm_args.verbosity_fpcm_move
    # 0) extract raw tensors as tuple
    tensors= tuple([a,env.C[env.keyC],env.T[env.keyT]])
    
    # function wrapping up the core of the FPCM MOVE segment of CTM algorithm
    def fpcm_MOVE_sl_c(*tensors):
        a, C, T= tensors

        # 1) compute fixed points for U, T
        i=0
        e0=1.0e+16
        tmp_=[(T,C)]
        while e0 > ctm_args.fpcm_fpt_tol:
            if verbosity>0: log.info(f"fpcm_MOVE_sl_c iteration {i}")
            # 1.1) get C, U by solving pulling-through eq
            #      --Cp--T-- = --U--Cp-- where 0--U--1
            #            |       |                2
            Cp, U= isogauge_MPS(T, C0=tmp_[-1][1], isogauge_tol=ctm_args.fpcm_isogauge_tol,\
                verbosity=verbosity)
            Tp= fp_T(a, U, T0=tmp_[-1][0])
            
            # (optional) check symmetry error on Tp
            Tp= 0.5*(Tp + Tp.permute(1,0,2))
            
            # 1.2) check convergence 
            tmp_.append((Tp,Cp))
            if len(tmp_)>1:
                e0= torch.linalg.norm(tmp_[1][0]-tmp_[0][0])/max(Tp.size())
                if verbosity>0: log.info(f"fpcm_MOVE_sl_c iteration {i} error (p=2) {e0}"\
                    +f" (p=inf) {torch.linalg.norm(tmp_[1][0]-tmp_[0][0],float('inf'))}")
                tmp_.pop(0)
            i+=1

        # 2) compute fixed point of the corner
        nC= fp_C(a, Tp, U, verbosity=verbosity)
        
        # 4) symmetrize, normalize and assign new C,T
        nC= 0.5*(nC + nC.t())
        nC= nC/torch.max(torch.abs(nC))
        # nC= nC/nC.norm()
        Tp= Tp/torch.max(torch.abs(Tp))
        # nT= nT/nT.norm()

        return nC, Tp

    # Call the core function, allowing for checkpointing
    if ctm_args.fwd_checkpoint_move:
        new_tensors= checkpoint(fpcm_MOVE_sl_c,*tensors)
    else:
        new_tensors= fpcm_MOVE_sl_c(*tensors)

    env.C[env.keyC]= new_tensors[0]
    env.T[env.keyT]= new_tensors[1]
//  [python]
def _l2normalize(v, eps=1e-10):
    return v / (torch.norm(v,dim=2,keepdim=True) + eps)
//  [python]
def _l2normalize(v, eps=1e-10):
    return v / (torch.linalg.norm(v,dim=2,keepdim=True) + eps)
//  [python]
def isogauge_MPS(T, C0=None, isogauge_tol=1.0e-8, verbosity=0):
    # get the fixed-point of MPS's transfer matrix
    #
    # C^2--T--           C^2--
    # |    |   = \lambda |
    # \----T--           \----
    if C0 is not None:
        C2_0= C0@C0
        C2_0= C2_0.view(-1)
    else:
        C2_0= None

    nC2= fp_TT(T, C2_0=C2_0, verbosity=verbosity)

    # (optional) verify hermicity
    nC2_asymm_norm= torch.norm(nC2-nC2.t())
    # print(f"nC2_asymm {nC2_asymm_norm}")
    assert( nC2_asymm_norm/torch.abs(nC2).max() < 1.0e-8 )
    nC2= 0.5*(nC2+nC2.t())
    
    D, U= torch.symeig(nC2, eigenvectors=True)
    # D might have an overall minus sign
    # NOTE: torch.symeig orders eigenpairs in ascending manner
    # order in descending order by magnitude
    absD,p= torch.sort(torch.abs(D),descending=True)
    D=D[p]
    U=U[:,p]
    if D[0]<0: D=-D
    # (optional) verify positivity
    assert D.min()/D[0]>-1.0e-14, "Fixed point of transfer matrix is not positive"
    # no large negative elements, clamp non-negative numbers
    D= torch.clamp(D,min=0)

    # take square root and normalize 
    # 
    # C^2--   C--
    # |     = |
    # \----   C--
    nC= U @ torch.diag(torch.sqrt(D/D[0])) @ U.t()

    # solve initial pulling-through equation
    # --C--T-- \propto --U--C-- where 0--U--1
    #      |             |               2
    P, U= pull_through(nC,T)

    # initial error in gauging (Frobenius or spectral norm ?)
    e0= torch.norm(nC-P)/max(nC.size())
    if verbosity>0: 
        log.info(f"isogauge_CT init gauging error (p=2) {e0} (p=inf) {torch.norm(nC-P,float('inf'))}")

    # is error above tolerance ? Iteratively refine
    while e0 > isogauge_tol:
        # get fixed point of mixed transfer matrix
        #
        # C--T--           C--
        # |  |   = \lambda |
        # \--U--           \--
        nC= fp_TT(T, U=U, C2_0=nC, verbosity=verbosity)  
        
        # left polar decomposition nC = QP and nC<-P where 
        # P is positive and Hermitian and Q is unitary 
        nC, Q= polar_decomp_left(nC, normalize=True)

        # (optional) verify nC is positive and Hermitian
        # solve pulling-through equation 
        # --nC--T-- \propto --U--nC--
        #       |                 |
        P, U= pull_through(nC, T)
        e0= torch.norm(nC-P)
        if verbosity>0: log.info(f"isogauge_CT iterative improvement gauging error (p=2) {e0}"\
            +f" (p=inf) {torch.norm(nC-P,float('inf'))}")
    return nC, U
//  [python]
def isogauge_MPS(T, C0=None, isogauge_tol=1.0e-8, verbosity=0):
    # get the fixed-point of MPS's transfer matrix
    #
    # C^2--T--           C^2--
    # |    |   = \lambda |
    # \----T--           \----
    if C0 is not None:
        C2_0= C0@C0
        C2_0= C2_0.view(-1)
    else:
        C2_0= None

    nC2= fp_TT(T, C2_0=C2_0, verbosity=verbosity)

    # (optional) verify hermicity
    nC2_asymm_norm= torch.linalg.norm(nC2-nC2.t())
    # print(f"nC2_asymm {nC2_asymm_norm}")
    assert( nC2_asymm_norm/torch.abs(nC2).max() < 1.0e-8 )
    nC2= 0.5*(nC2+nC2.t())
    
    D, U= torch.symeig(nC2, eigenvectors=True)
    # D might have an overall minus sign
    # NOTE: torch.symeig orders eigenpairs in ascending manner
    # order in descending order by magnitude
    absD,p= torch.sort(torch.abs(D),descending=True)
    D=D[p]
    U=U[:,p]
    if D[0]<0: D=-D
    # (optional) verify positivity
    assert D.min()/D[0]>-1.0e-14, "Fixed point of transfer matrix is not positive"
    # no large negative elements, clamp non-negative numbers
    D= torch.clamp(D,min=0)

    # take square root and normalize 
    # 
    # C^2--   C--
    # |     = |
    # \----   C--
    nC= U @ torch.diag(torch.sqrt(D/D[0])) @ U.t()

    # solve initial pulling-through equation
    # --C--T-- \propto --U--C-- where 0--U--1
    #      |             |               2
    P, U= pull_through(nC,T)

    # initial error in gauging (Frobenius or spectral norm ?)
    e0= torch.linalg.norm(nC-P)/max(nC.size())
    if verbosity>0: 
        log.info(f"isogauge_CT init gauging error (p=2) {e0} (p=inf) {torch.linalg.norm(nC-P,float('inf'))}")

    # is error above tolerance ? Iteratively refine
    while e0 > isogauge_tol:
        # get fixed point of mixed transfer matrix
        #
        # C--T--           C--
        # |  |   = \lambda |
        # \--U--           \--
        nC= fp_TT(T, U=U, C2_0=nC, verbosity=verbosity)  
        
        # left polar decomposition nC = QP and nC<-P where 
        # P is positive and Hermitian and Q is unitary 
        nC, Q= polar_decomp_left(nC, normalize=True)

        # (optional) verify nC is positive and Hermitian
        # solve pulling-through equation 
        # --nC--T-- \propto --U--nC--
        #       |                 |
        P, U= pull_through(nC, T)
        e0= torch.linalg.norm(nC-P)
        if verbosity>0: log.info(f"isogauge_CT iterative improvement gauging error (p=2) {e0}"\
            +f" (p=inf) {torch.linalg.norm(nC-P,float('inf'))}")
    return nC, U
//  [python]
def update_weight_Boosting(self, weight_mat, dist_old, dist_new):
        epsilon = 1e-12
        dist_old = dist_old.detach()
        dist_new = dist_new.detach()
        ind = dist_new > dist_old + epsilon
        weight_mat[ind] = weight_mat[ind] * \
            (1 + torch.sigmoid(dist_new[ind] - dist_old[ind]))
        weight_norm = torch.norm(weight_mat, dim=1, p=1)
        weight_mat = weight_mat / weight_norm.t().unsqueeze(1).repeat(1, self.len_seq)
        return weight_mat
//  [python]
def update_weight_Boosting(self, weight_mat, dist_old, dist_new):
        epsilon = 1e-12
        dist_old = dist_old.detach()
        dist_new = dist_new.detach()
        ind = dist_new > dist_old + epsilon
        weight_mat[ind] = weight_mat[ind] * \
            (1 + torch.sigmoid(dist_new[ind] - dist_old[ind]))
        weight_norm = torch.linalg.norm(weight_mat, dim=1, ord=1)
        weight_mat = weight_mat / weight_norm.t().unsqueeze(1).repeat(1, self.len_seq)
        return weight_mat
//  [python]
def rotate_a_b_axis_angle_torch_batched(a, b):
    a = a / torch.norm(a, dim=1, keepdim=True)
    b = b / torch.norm(b, dim=1, keepdim=True)
    rot_axis = torch.cross(a, b)

    a_proj = b * torch.sum(a * b, dim=1, keepdim=True)
    a_ort = a - a_proj
    theta = np.arctan2(
        torch.norm(a_ort, dim=1, keepdim=True),
        torch.norm(a_proj, dim=1, keepdim=True)
    )

    theta[torch.sum(a * b, dim=1) < 0] = np.pi - theta[torch.sum(a * b, dim=1) < 0]

    aa = rot_axis / torch.norm(rot_axis, dim=1, keepdim=True) * theta

    return aa
//  [python]
def rotate_a_b_axis_angle_torch_batched(a, b):
    a = a / torch.linalg.norm(a, dim=1, keepdim=True)
    b = b / torch.linalg.norm(b, dim=1, keepdim=True)
    rot_axis = torch.cross(a, b)

    a_proj = b * torch.sum(a * b, dim=1, keepdim=True)
    a_ort = a - a_proj
    theta = np.arctan2(
        torch.linalg.norm(a_ort, dim=1, keepdim=True),
        torch.linalg.norm(a_proj, dim=1, keepdim=True)
    )

    theta[torch.sum(a * b, dim=1) < 0] = np.pi - theta[torch.sum(a * b, dim=1) < 0]

    aa = rot_axis / torch.linalg.norm(rot_axis, dim=1, keepdim=True) * theta

    return aa
//  [python]
def update_C(self, zs: torch.Tensor, ys: torch.Tensor, assigned_weights: torch.Tensor, h_sig: torch.Tensor) -> None:
        """Update the covariance shape matrix C based on rank-1 and rank-mu updates
        This operation is bounded O(d^2 popsize), which is associated with computing the rank-mu update (summing across popsize d*d matrices)
        Args:
            zs (torch.Tensor): A tensor of shape [popsize, d] of samples from the local coordinate space e.g. z_i ~ N(0, I_d)
            ys (torch.Tensor): A tensor of shape [popsize, d] of samples from the shaped coordinate space e.g. y_i ~ N(0, C)
            assigned_weights (torch.Tensor): A [popsize, ] dimensional tensor of ordered weights
            h_sig (torch.Tensor): Whether to stall the update based on the evolution path on sigma, p_sigma, expressed as a torch float
        """
        d = self._problem.solution_length
        # If using Active CMA-ES, reweight negative weights
        if self.active:
            assigned_weights = torch.where(
                assigned_weights > 0, assigned_weights, d * assigned_weights / torch.norm(zs, dim=-1).pow(2.0)
            )
        c1a = self.c_1 * (1 - (1 - h_sig**2) * self.c_c * (2 - self.c_c))  # adjust for variance loss
        weighted_pc = (self.c_1 / (c1a + 1e-23)) ** 0.5
        if self.separable:
            # Rank-1 update
            r1_update = c1a * (self.p_c.pow(2.0) - self.C)
            # Rank-mu update
            rmu_update = self.c_mu * torch.sum(
                assigned_weights.unsqueeze(-1) * (ys.pow(2.0) - self.C.unsqueeze(0)), dim=0
            )
        else:
            # Rank-1 update
            r1_update = c1a * (torch.outer(weighted_pc * self.p_c, weighted_pc * self.p_c) - self.C)
            # Rank-mu update
            rmu_update = self.c_mu * (
                torch.sum(assigned_weights.unsqueeze(-1).unsqueeze(-1) * (ys.unsqueeze(1) * ys.unsqueeze(2)), dim=0)
                - torch.sum(self.weights) * self.C
            )

        # Update C
        self.C = self.C + r1_update + rmu_update
//  [python]
def update_C(self, zs: torch.Tensor, ys: torch.Tensor, assigned_weights: torch.Tensor, h_sig: torch.Tensor) -> None:
        """Update the covariance shape matrix C based on rank-1 and rank-mu updates
        This operation is bounded O(d^2 popsize), which is associated with computing the rank-mu update (summing across popsize d*d matrices)
        Args:
            zs (torch.Tensor): A tensor of shape [popsize, d] of samples from the local coordinate space e.g. z_i ~ N(0, I_d)
            ys (torch.Tensor): A tensor of shape [popsize, d] of samples from the shaped coordinate space e.g. y_i ~ N(0, C)
            assigned_weights (torch.Tensor): A [popsize, ] dimensional tensor of ordered weights
            h_sig (torch.Tensor): Whether to stall the update based on the evolution path on sigma, p_sigma, expressed as a torch float
        """
        d = self._problem.solution_length
        # If using Active CMA-ES, reweight negative weights
        if self.active:
            assigned_weights = torch.where(
                assigned_weights > 0, assigned_weights, d * assigned_weights / torch.linalg.norm(zs, dim=-1).pow(2.0)
            )
        c1a = self.c_1 * (1 - (1 - h_sig**2) * self.c_c * (2 - self.c_c))  # adjust for variance loss
        weighted_pc = (self.c_1 / (c1a + 1e-23)) ** 0.5
        if self.separable:
            # Rank-1 update
            r1_update = c1a * (self.p_c.pow(2.0) - self.C)
            # Rank-mu update
            rmu_update = self.c_mu * torch.sum(
                assigned_weights.unsqueeze(-1) * (ys.pow(2.0) - self.C.unsqueeze(0)), dim=0
            )
        else:
            # Rank-1 update
            r1_update = c1a * (torch.outer(weighted_pc * self.p_c, weighted_pc * self.p_c) - self.C)
            # Rank-mu update
            rmu_update = self.c_mu * (
                torch.sum(assigned_weights.unsqueeze(-1).unsqueeze(-1) * (ys.unsqueeze(1) * ys.unsqueeze(2)), dim=0)
                - torch.sum(self.weights) * self.C
            )

        # Update C
        self.C = self.C + r1_update + rmu_update
//  [python]
def rodrigues(theta):
    """Convert axis-angle representation to rotation matrix.
    Args:
        theta: size = [B, 3]
    Returns:
        Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
    """
    l1norm = torch.norm(theta + 1e-8, p = 2, dim = 1)
    angle = torch.unsqueeze(l1norm, -1)
    normalized = torch.div(theta, angle)
    angle = angle * 0.5
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    quat = torch.cat([v_cos, v_sin * normalized], dim = 1)
    return quat2mat(quat)
//  [python]
def rodrigues(theta):
    """Convert axis-angle representation to rotation matrix.
    Args:
        theta: size = [B, 3]
    Returns:
        Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
    """
    l1norm = torch.linalg.norm(theta + 1e-8, ord=2, dim=1)
    angle = torch.unsqueeze(l1norm, -1)
    normalized = torch.div(theta, angle)
    angle = angle * 0.5
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    quat = torch.cat([v_cos, v_sin * normalized], dim=1)
    return quat2mat(quat)
//  [python]
def contrastive_loss(self, F0_p, F1_p, F1, pts1):
        dist_pos = torch.norm(F0_p - F1_p, p=2, dim=-1)
        loss_pos = torch.clamp(dist_pos-self.conf.loss.margin_positive, min=0)
        loss_pos = loss_pos**2

        all_dist_neg = distance_s2d(F0_p.detach(), F1.detach())

        c, h, w = F1.shape[-3:]
        if self.conf.loss.windowed_negative_sampling:
            radius = self.conf.loss.ratio_radius_negative * max(h, w)
            mask_x = torch.abs(torch.arange(w, device=pts1.device)[None, None]
                               - pts1[..., :1]) < radius
            mask_y = torch.abs(torch.arange(h, device=pts1.device)[None, None]
                               - pts1[..., 1:]) < radius
            mask_neg = mask_x[..., None, :] & mask_y[..., None]
            all_dist_neg.masked_fill_(mask_neg, float('inf'))
        all_dist_neg = all_dist_neg.reshape(all_dist_neg.shape[:-2]+(-1,))

        n = self.conf.loss.num_top_negative_sampling
        if n > 0:
            inds = torch.topk(all_dist_neg, n, dim=-1, largest=False).indices
            sampled = torch.randint_like(inds[..., 0], high=n)
            inds = inds.gather(-1, sampled[..., None]).squeeze(-1)
        else:
            n = all_dist_neg.shape[-1]
            inds = torch.randint(n, all_dist_neg.shape[:-1])
            inds = inds.to(all_dist_neg.device)

        # For an unknown reason, gather generates huge grad tensor,
        # resulting in OOM, while indexing works just fine.
        inds = torch.stack([inds // w, inds % w], -1)  # y, x
        F1_neg = torch.stack([F.permute(1, 2, 0)[tuple(i.T)]
                              for F, i in zip(F1, inds)], 0)
        # F1_neg = F1.reshape(-1, 1, c, h*w).expand(inds.shape + (-1, -1))
        # inds = inds[..., None, None].expand(-1, -1, c, 1)
        # F1_neg = F1_neg.gather(-1, inds).squeeze(-1)

        # Different formulation in GN-Net vs LM-Reloc: square after or before.
        # Here we pick the formulation of LM-Reloc.
        dist_neg = torch.sum((F0_p - F1_neg)**2, -1)
        # dist_neg = torch.norm(F0_p - F1_neg, p=2, dim=-1)
        loss_neg = torch.clamp(self.conf.loss.margin_negative-dist_neg, min=0)
        # loss_neg = loss_neg**2
        return loss_pos + loss_neg
//  [python]
def contrastive_loss(self, F0_p, F1_p, F1, pts1):
        dist_pos = torch.linalg.norm(F0_p - F1_p, ord=2, dim=-1)
        loss_pos = torch.clamp(dist_pos-self.conf.loss.margin_positive, min=0)
        loss_pos = loss_pos**2

        all_dist_neg = distance_s2d(F0_p.detach(), F1.detach())

        c, h, w = F1.shape[-3:]
        if self.conf.loss.windowed_negative_sampling:
            radius = self.conf.loss.ratio_radius_negative * max(h, w)
            mask_x = torch.abs(torch.arange(w, device=pts1.device)[None, None]
                               - pts1[..., :1]) < radius
            mask_y = torch.abs(torch.arange(h, device=pts1.device)[None, None]
                               - pts1[..., 1:]) < radius
            mask_neg = mask_x[..., None, :] & mask_y[..., None]
            all_dist_neg.masked_fill_(mask_neg, float('inf'))
        all_dist_neg = all_dist_neg.reshape(all_dist_neg.shape[:-2]+(-1,))

        n = self.conf.loss.num_top_negative_sampling
        if n > 0:
            inds = torch.topk(all_dist_neg, n, dim=-1, largest=False).indices
            sampled = torch.randint_like(inds[..., 0], high=n)
            inds = inds.gather(-1, sampled[..., None]).squeeze(-1)
        else:
            n = all_dist_neg.shape[-1]
            inds = torch.randint(n, all_dist_neg.shape[:-1])
            inds = inds.to(all_dist_neg.device)

        # For an unknown reason, gather generates huge grad tensor,
        # resulting in OOM, while indexing works just fine.
        inds = torch.stack([inds // w, inds % w], -1)  # y, x
        F1_neg = torch.stack([F.permute(1, 2, 0)[tuple(i.T)]
                              for F, i in zip(F1, inds)], 0)
        # F1_neg = F1.reshape(-1, 1, c, h*w).expand(inds.shape + (-1, -1))
        # inds = inds[..., None, None].expand(-1, -1, c, 1)
        # F1_neg = F1_neg.gather(-1, inds).squeeze(-1)

        # Different formulation in GN-Net vs LM-Reloc: square after or before.
        # Here we pick the formulation of LM-Reloc.
        dist_neg = torch.sum((F0_p - F1_neg)**2, -1)
        # dist_neg = torch.linalg.norm(F0_p - F1_neg, ord=2, dim=-1)
        loss_neg = torch.clamp(self.conf.loss.margin_negative-dist_neg, min=0)
        # loss_neg = loss_neg**2
        return loss_pos + loss_neg
//  [python]
def reconstruction_distance(D, cur_Z, last_Z):
    distance = torch.norm(D.mm(last_Z - cur_Z), p=2, dim=0) / torch.norm(D.mm(last_Z), p=2, dim=0)
    max_distance = distance.max()
    return distance, max_distance
//  [python]
def reconstruction_distance(D, cur_Z, last_Z):
    distance = torch.linalg.norm(D.mm(last_Z - cur_Z), ord=2, dim=0) / torch.linalg.norm(D.mm(last_Z), ord=2, dim=0)
    max_distance = distance.max()
    return distance, max_distance
//  [python]
def metrics(self, pred, data):
        p3D_ref = data['ref']['points3D']
        T_init = data['T_r2q_init']
        T_r2q_gt = data['T_r2q_gt']
        T_q2r_gt = T_r2q_gt.inv()

        T_r2q_init = []
        T_r2q_opt = []
        for i in reversed(range(len(self.extractor.scales))):
            F_ref = pred['ref']['feature_maps'][i]
            F_q = pred['query']['feature_maps'][i]
            cam_ref = pred['ref']['camera_pyr'][i]
            cam_q = pred['query']['camera_pyr'][i]

            p2D_ref, visible = cam_ref.world2image(p3D_ref)
            F_ref, mask, _ = self.optimizer.interpolator(F_ref, p2D_ref)
            mask &= visible

            if self.conf.normalize_features:
                F_ref = nnF.normalize(F_ref, dim=2)
                F_q = nnF.normalize(F_q, dim=1)

            T_opt, failed = self.optimizer(dict(
                p3D=p3D_ref, F_ref=F_ref, F_q=F_q, T_init=T_init, cam_q=cam_q,
                mask=mask))

            T_r2q_init.append(T_init)
            T_r2q_opt.append(T_opt)
            T_init = T_opt.detach()

        def scaled_pose_error(T_r2q):
            err_R, err_t = (T_r2q @ T_q2r_gt).magnitude()
            err_t /= torch.norm(T_q2r_gt.t, dim=-1)
            return err_R, err_t

        metrics = {}
        # Compute pose errors
        for i, T_opt in enumerate(T_r2q_opt):
            err = scaled_pose_error(T_opt)
            metrics[f'R_error/{i}'], metrics[f't_error/{i}'] = err
        metrics['R_error'], metrics['t_error'] = err
        err_init = scaled_pose_error(T_r2q_init[0])
        metrics['R_error/init'], metrics['t_error/init'] = err_init

        # Compute reprojection errors
        def project(T_r2q):
            return cam_q.world2image(T_r2q * p3D_ref)

        p2D_q_gt, mask = project(T_r2q_gt)
        p2D_q_i, mask_i = project(T_r2q_init[0])

        def reprojection_error(T_r2q):
            p2D_q, _ = project(T_r2q)
            err = torch.sum((p2D_q_gt - p2D_q)**2, dim=-1)
            err = scaled_barron(1., 2.)(err)[0]/4
            err = masked_mean(err, mask, -1)
            return err

        for i, T_opt in enumerate(T_r2q_opt):
            err = reprojection_error(T_opt).clamp(max=50)
            metrics[f'loss/reprojection_error/{i}'] = err
        metrics['loss/reprojection_error'] = err

        err_init = reprojection_error(T_r2q_init[0])
        metrics['loss/reprojection_error/init'] = err_init

        return metrics
//  [python]
def metrics(self, pred, data):
        p3D_ref = data['ref']['points3D']
        T_init = data['T_r2q_init']
        T_r2q_gt = data['T_r2q_gt']
        T_q2r_gt = T_r2q_gt.inv()

        T_r2q_init = []
        T_r2q_opt = []
        for i in reversed(range(len(self.extractor.scales))):
            F_ref = pred['ref']['feature_maps'][i]
            F_q = pred['query']['feature_maps'][i]
            cam_ref = pred['ref']['camera_pyr'][i]
            cam_q = pred['query']['camera_pyr'][i]

            p2D_ref, visible = cam_ref.world2image(p3D_ref)
            F_ref, mask, _ = self.optimizer.interpolator(F_ref, p2D_ref)
            mask &= visible

            if self.conf.normalize_features:
                F_ref = nnF.normalize(F_ref, dim=2)
                F_q = nnF.normalize(F_q, dim=1)

            T_opt, failed = self.optimizer(dict(
                p3D=p3D_ref, F_ref=F_ref, F_q=F_q, T_init=T_init, cam_q=cam_q,
                mask=mask))

            T_r2q_init.append(T_init)
            T_r2q_opt.append(T_opt)
            T_init = T_opt.detach()

        def scaled_pose_error(T_r2q):
            err_R, err_t = (T_r2q @ T_q2r_gt).magnitude()
            err_t /= torch.linalg.norm(T_q2r_gt.t, dim=-1)
            return err_R, err_t

        metrics = {}
        # Compute pose errors
        for i, T_opt in enumerate(T_r2q_opt):
            err = scaled_pose_error(T_opt)
            metrics[f'R_error/{i}'], metrics[f't_error/{i}'] = err
        metrics['R_error'], metrics['t_error'] = err
        err_init = scaled_pose_error(T_r2q_init[0])
        metrics['R_error/init'], metrics['t_error/init'] = err_init

        # Compute reprojection errors
        def project(T_r2q):
            return cam_q.world2image(T_r2q * p3D_ref)

        p2D_q_gt, mask = project(T_r2q_gt)
        p2D_q_i, mask_i = project(T_r2q_init[0])

        def reprojection_error(T_r2q):
            p2D_q, _ = project(T_r2q)
            err = torch.sum((p2D_q_gt - p2D_q)**2, dim=-1)
            err = scaled_barron(1., 2.)(err)[0]/4
            err = masked_mean(err, mask, -1)
            return err

        for i, T_opt in enumerate(T_r2q_opt):
            err = reprojection_error(T_opt).clamp(max=50)
            metrics[f'loss/reprojection_error/{i}'] = err
        metrics['loss/reprojection_error'] = err

        err_init = reprojection_error(T_r2q_init[0])
        metrics['loss/reprojection_error/init'] = err_init

        return metrics
//  [python]
def stable_norm(input, *args, **kwargs):
    return torch.norm(input, *args, **kwargs)
    input = input.clone()
    with torch.no_grad():
        sign = torch.sign(input)
        input = torch.abs(input)
        input.clamp_(min=1e-10)
        input = sign * input
    return torch.norm(input, *args, **kwargs)
//  [python]
def stable_norm(input, *args, **kwargs):
    return torch.linalg.norm(input, *args, **kwargs)
    input = input.clone()
    with torch.no_grad():
        sign = torch.sign(input)
        input = torch.abs(input)
        input.clamp_(min=1e-10)
        input = sign * input
    return torch.linalg.norm(input, *args, **kwargs)
//  [python]
def robust_PCA():
    d1 = 3
    d2 = 4
    torch.manual_seed(1)
    eta = .05
    Y = torch.randn(d1,d2).to(device=device, dtype=torch_dtype)

    # variables and corresponding dimensions.
    var_in = {"M": [d1,d2],"S": [d1,d2]}


    opts = pygransoStruct()
    opts.print_frequency = 10
    opts.x0 = .2 * torch.ones((2*d1*d2,1)).to(device=device, dtype=torch_dtype)
    opts.opt_tol = 1e-6
    opts.maxit = 50
    opts.print_level = print_level
    opts.double_precision = double_precision
    opts.torch_device = device

    def comb_fn(X_struct):
        M = X_struct.M
        S = X_struct.S

        # objective function
        f = torch.norm(M, p = 'nuc') + eta * torch.norm(S, p = 1)

        # inequality constraint, matrix form
        ci = None

        # equality constraint
        ce = pygransoStruct()
        ce.c1 = M + S - Y

        return [f,ci,ce]

    soln = pygranso(var_spec= var_in, combined_fn = comb_fn, user_opts = opts)
    print("test 4/7 passed (Robust PCA)")
//  [python]
def robust_PCA():
    d1 = 3
    d2 = 4
    torch.manual_seed(1)
    eta = .05
    Y = torch.randn(d1,d2).to(device=device, dtype=torch_dtype)

    # variables and corresponding dimensions.
    var_in = {"M": [d1,d2],"S": [d1,d2]}


    opts = pygransoStruct()
    opts.print_frequency = 10
    opts.x0 = .2 * torch.ones((2*d1*d2,1)).to(device=device, dtype=torch_dtype)
    opts.opt_tol = 1e-6
    opts.maxit = 50
    opts.print_level = print_level
    opts.double_precision = double_precision
    opts.torch_device = device

    def comb_fn(X_struct):
        M = X_struct.M
        S = X_struct.S

        # objective function
        f = torch.linalg.norm(M, ord='nuc') + eta * torch.linalg.norm(S, ord=1)

        # inequality constraint, matrix form
        ci = None

        # equality constraint
        ce = pygransoStruct()
        ce.c1 = M + S - Y

        return [f,ci,ce]

    soln = pygranso(var_spec= var_in, combined_fn = comb_fn, user_opts = opts)
    print("test 4/7 passed (Robust PCA)")
//  [python]
def lasso():
    n = 80
    eta = 0.5 # parameter for penalty term
    torch.manual_seed(1)
    b = torch.rand(n,1)
    pos_one = torch.ones(n-1)
    neg_one = -torch.ones(n-1)
    F = torch.zeros(n-1,n)
    F[:,0:n-1] += torch.diag(neg_one,0)
    F[:,1:n] += torch.diag(pos_one,0)
    F = F.to(device=device, dtype=torch_dtype)  # double precision requireed in torch operations
    b = b.to(device=device, dtype=torch_dtype)

    # variables and corresponding dimensions.
    var_in = {"x": [n,1]}


    def comb_fn(X_struct):
        x = X_struct.x

        # objective function
        f = (x-b).t() @ (x-b)  + eta * torch.norm( F@x, p = 1)

        # inequality constraint, matrix form
        ci = None
        # equality constraint
        ce = None

        return [f,ci,ce]

    opts = pygransoStruct()
    opts.QPsolver = 'osqp'
    opts.x0 = torch.ones((n,1)).to(device=device, dtype=torch_dtype)
    opts.print_level = 1
    opts.print_frequency = 10
    opts.print_level = print_level
    opts.double_precision = double_precision
    opts.maxit = 30
    opts.torch_device = device

    soln = pygranso(var_spec= var_in, combined_fn = comb_fn, user_opts = opts)

    print("test 5/7 passed (LASSO)")
//  [python]
def lasso():
    n = 80
    eta = 0.5 # parameter for penalty term
    torch.manual_seed(1)
    b = torch.rand(n,1)
    pos_one = torch.ones(n-1)
    neg_one = -torch.ones(n-1)
    F = torch.zeros(n-1,n)
    F[:,0:n-1] += torch.diag(neg_one,0)
    F[:,1:n] += torch.diag(pos_one,0)
    F = F.to(device=device, dtype=torch_dtype)  # double precision requireed in torch operations
    b = b.to(device=device, dtype=torch_dtype)

    # variables and corresponding dimensions.
    var_in = {"x": [n,1]}


    def comb_fn(X_struct):
        x = X_struct.x

        # objective function
        f = (x-b).t() @ (x-b)  + eta * torch.linalg.norm( F@x, ord = 1)

        # inequality constraint, matrix form
        ci = None
        # equality constraint
        ce = None

        return [f,ci,ce]

    opts = pygransoStruct()
    opts.QPsolver = 'osqp'
    opts.x0 = torch.ones((n,1)).to(device=device, dtype=torch_dtype)
    opts.print_level = 1
    opts.print_frequency = 10
    opts.print_level = print_level
    opts.double_precision = double_precision
    opts.maxit = 30
    opts.torch_device = device

    soln = pygranso(var_spec= var_in, combined_fn = comb_fn, user_opts = opts)

    print("test 5/7 passed (LASSO)")
//  [python]
def norm(self, X, G):
        """
        Norm on the tangent space of the Stiefel is simply the Euclidean
        norm.
        """
        return torch.norm(G)
//  [python]
def norm(self, X, G):
        """
        Norm on the tangent space of the Stiefel is simply the Euclidean
        norm.
        """
        return torch.linalg.norm(G)
//  [python]
def norm(self, p="fro", dim=None, keepdim=False, dtype=None):
        r"""See :func:`torch.norm`"""
        return torch.norm(self, p, dim, keepdim, dtype=dtype)
//  [python]
def norm(self, p="fro", dim=None, keepdim=False, dtype=None):
        r"""See :func:`torch.norm`"""
        return torch.linalg.norm(self, p, dim, keepdim, dtype=dtype)
//  [python]
def RBF(x):
    d_total_t = torch.from_numpy(d_total).unsqueeze(1).cuda()
    w_t = torch.from_numpy(w).cuda()
    x_l = x.unsqueeze(0).repeat(len(d_total_t), 1, 1) # 创立一个足够大的x矩阵
    # 得到R大矩阵
    
    R = torch.norm(d_total_t - x_l, dim=2)
    #Rn = -(x_l[:, :, 0]-d_total_t[:, :, 0]) # 试一试我自己创建的theta的径向基函数
    y = torch.mm(torch.exp(-gama*R.T), w_t)
    #y = torch.mm(torch.sqrt(0.5*(R.T-Rn.T)), w_t)# 试一试我自己创建的theta的径向基函数
    #y = torch.mm(torch.sqrt(R.T), w_t)
    return y
//  [python]
def RBF(x):
    d_total_t = torch.from_numpy(d_total).unsqueeze(1).cuda()
    w_t = torch.from_numpy(w).cuda()
    x_l = x.unsqueeze(0).repeat(len(d_total_t), 1, 1) # 创立一个足够大的x矩阵
    # 得到R大矩阵
    
    R = torch.linalg.norm(d_total_t - x_l, dim=2)
    #Rn = -(x_l[:, :, 0]-d_total_t[:, :, 0]) # 试一试我自己创建的theta的径向基函数
    y = torch.mm(torch.exp(-gama*R.T), w_t)
    #y = torch.mm(torch.sqrt(0.5*(R.T-Rn.T)), w_t)# 试一试我自己创建的theta的径向基函数
    #y = torch.mm(torch.sqrt(R.T), w_t)
    return y
//  [python]
def evaluate():
    N_test = 100
    dom_koch_n = koch_points.get_koch_points_lin(N_test) # 均匀在科赫雪花内部步点
    dom_koch_t = torch.tensor(dom_koch_n, device = 'cuda')
    
    u_pred = pred(dom_koch_t)
    u_pred = u_pred.data.cpu()
    

    u_exact = np.zeros((len(dom_koch_n) ,1))
    u_exact[(dom_koch_n[:,0]**2+dom_koch_n[:,1]**2)<r0**2] = 1/a1*np.linalg.norm(dom_koch_n[(dom_koch_n[:,0]**2+dom_koch_n[:,1]**2)<r0**2], axis=1, keepdims=True)**4 # 现在网格里面赋值
    u_exact[(dom_koch_n[:,0]**2+dom_koch_n[:,1]**2)>=r0**2] = 1/a2*np.linalg.norm(dom_koch_n[(dom_koch_n[:,0]**2+dom_koch_n[:,1]**2)>=r0**2], axis=1, keepdims=True)**4 + r0**4*(1/a1-1/a2) # 现在网格外面赋值

    u_exact = torch.from_numpy(u_exact)
    error = torch.abs(u_pred - u_exact) # get the error in every points
    error_t = torch.norm(error)/torch.norm(u_exact) # get the total relative L2 error
    return error_t
//  [python]
def evaluate():
    N_test = 100
    dom_koch_n = koch_points.get_koch_points_lin(N_test) # 均匀在科赫雪花内部步点
    dom_koch_t = torch.tensor(dom_koch_n, device = 'cuda')
    
    u_pred = pred(dom_koch_t)
    u_pred = u_pred.data.cpu()
    

    u_exact = np.zeros((len(dom_koch_n) ,1))
    u_exact[(dom_koch_n[:,0]**2+dom_koch_n[:,1]**2)<r0**2] = 1/a1*np.linalg.norm(dom_koch_n[(dom_koch_n[:,0]**2+dom_koch_n[:,1]**2)<r0**2], axis=1, keepdims=True)**4 # 现在网格里面赋值
    u_exact[(dom_koch_n[:,0]**2+dom_koch_n[:,1]**2)>=r0**2] = 1/a2*np.linalg.norm(dom_koch_n[(dom_koch_n[:,0]**2+dom_koch_n[:,1]**2)>=r0**2], axis=1, keepdims=True)**4 + r0**4*(1/a1-1/a2) # 现在网格外面赋值

    u_exact = torch.from_numpy(u_exact)
    error = torch.abs(u_pred - u_exact) # get the error in every points
    error_t = torch.linalg.norm(error)/torch.linalg.norm(u_exact) # get the total relative L2 error
    return error_t
//  [python]
def create_cost_and_derivates(manifold, matrix, backend):
    euclidean_gradient = None

    if backend == "autograd":

        @pymanopt.function.autograd(manifold)
        def cost(u, s, vt):
            X = u @ np.diag(s) @ vt
            return np.linalg.norm(X - matrix) ** 2

    elif backend == "jax":

        @pymanopt.function.jax(manifold)
        def cost(u, s, vt):
            X = u @ jnp.diag(s) @ vt
            return jnp.linalg.norm(X - matrix) ** 2

    elif backend == "numpy":

        @pymanopt.function.numpy(manifold)
        def cost(u, s, vt):
            X = u @ np.diag(s) @ vt
            return np.linalg.norm(X - matrix) ** 2

        @pymanopt.function.numpy(manifold)
        def euclidean_gradient(u, s, vt):
            X = u @ np.diag(s) @ vt
            S = np.diag(s)
            gu = 2 * (X - matrix) @ (S @ vt).T
            gs = 2 * np.diag(u.T @ (X - matrix) @ vt.T)
            gvt = 2 * (u @ S).T @ (X - matrix)
            return gu, gs, gvt

    elif backend == "pytorch":
        matrix_ = torch.from_numpy(matrix)

        @pymanopt.function.pytorch(manifold)
        def cost(u, s, vt):
            X = u @ torch.diag(s) @ vt
            return torch.norm(X - matrix_) ** 2

    elif backend == "tensorflow":

        @pymanopt.function.tensorflow(manifold)
        def cost(u, s, vt):
            X = u @ tf.linalg.diag(s) @ vt
            return tf.norm(X - matrix) ** 2

    else:
        raise ValueError(f"Unsupported backend '{backend}'")

    return cost, euclidean_gradient
//  [python]
def create_cost_and_derivates(manifold, matrix, backend):
    euclidean_gradient = None

    if backend == "autograd":

        @pymanopt.function.autograd(manifold)
        def cost(u, s, vt):
            X = u @ np.diag(s) @ vt
            return np.linalg.norm(X - matrix) ** 2

    elif backend == "jax":

        @pymanopt.function.jax(manifold)
        def cost(u, s, vt):
            X = u @ jnp.diag(s) @ vt
            return jnp.linalg.norm(X - matrix) ** 2

    elif backend == "numpy":

        @pymanopt.function.numpy(manifold)
        def cost(u, s, vt):
            X = u @ np.diag(s) @ vt
            return np.linalg.norm(X - matrix) ** 2

        @pymanopt.function.numpy(manifold)
        def euclidean_gradient(u, s, vt):
            X = u @ np.diag(s) @ vt
            S = np.diag(s)
            gu = 2 * (X - matrix) @ (S @ vt).T
            gs = 2 * np.diag(u.T @ (X - matrix) @ vt.T)
            gvt = 2 * (u @ S).T @ (X - matrix)
            return gu, gs, gvt

    elif backend == "pytorch":
        matrix_ = torch.from_numpy(matrix)

        @pymanopt.function.pytorch(manifold)
        def cost(u, s, vt):
            X = u @ torch.diag(s) @ vt
            return torch.linalg.norm(X - matrix_) ** 2

    elif backend == "tensorflow":

        @pymanopt.function.tensorflow(manifold)
        def cost(u, s, vt):
            X = u @ tf.linalg.diag(s) @ vt
            return tf.norm(X - matrix) ** 2

    else:
        raise ValueError(f"Unsupported backend '{backend}'")

    return cost, euclidean_gradient
//  [python]
def distance_to_depth(distance: Float[Tensor, "B 1 H W"], fov_x: float, fov_y: float):
    device = distance.device
    batch_size, _, height, width = distance.shape
    fov_x_rad = fov_x * torch.pi / 180.0
    fov_y_rad = fov_y * torch.pi / 180.0
    fx = 1 / math.tan(fov_x_rad / 2)
    fy = 1 / math.tan(fov_y_rad / 2)
    x = torch.linspace(-1, 1, width, device=device) / fx
    y = torch.linspace(-1, 1, height, device=device) / fy
    ys, xs = torch.meshgrid([y, x], indexing="ij")
    xy = torch.stack([xs, ys]).unsqueeze(0).repeat(batch_size, 1, 1, 1)
    xyz = torch.cat([xy, torch.ones_like(xy[:, 0:1])], dim=1)
    xyz = xyz / torch.norm(xyz, dim=1, keepdim=True)
    xyz = xyz * distance
    depth = xyz[:, 2:3]
    return depth
//  [python]
def distance_to_depth(distance: Float[Tensor, "B 1 H W"], fov_x: float, fov_y: float):
    device = distance.device
    batch_size, _, height, width = distance.shape
    fov_x_rad = fov_x * torch.pi / 180.0
    fov_y_rad = fov_y * torch.pi / 180.0
    fx = 1 / math.tan(fov_x_rad / 2)
    fy = 1 / math.tan(fov_y_rad / 2)
    x = torch.linspace(-1, 1, width, device=device) / fx
    y = torch.linspace(-1, 1, height, device=device) / fy
    ys, xs = torch.meshgrid([y, x], indexing="ij")
    xy = torch.stack([xs, ys]).unsqueeze(0).repeat(batch_size, 1, 1, 1)
    xyz = torch.cat([xy, torch.ones_like(xy[:, 0:1])], dim=1)
    xyz = xyz / torch.linalg.norm(xyz, dim=1, keepdim=True)
    xyz = xyz * distance
    depth = xyz[:, 2:3]
    return depth
//  [python]
def multi_tensor_total_norm(grads, chunk_size=2048 * 32) -> torch.Tensor:
    per_device_grads = {}
    norms = []
    for grad in grads:
        device = grad.device
        cur_device_grads = per_device_grads.get(device)
        if cur_device_grads is None:
            cur_device_grads = []
            per_device_grads[device] = cur_device_grads
        cur_device_grads.append(grad)
    for device in per_device_grads.keys():
        cur_device_grads = per_device_grads[device]
        if device.type == "cuda":
            # TODO(msb) return has_inf
            has_inf = torch.zeros((1, 1), dtype=torch.int, device=device)
            with torch.cuda.device(device):
                norm = multi_tensor_l2norm(
                    chunk_size, has_inf, [cur_device_grads], False
                )
            norms.append(norm[0].to(torch.cuda.current_device()))
        else:
            norms += [torch.norm(g, p=2, dtype=torch.float32) for g in cur_device_grads]
    total_norm = torch.norm(torch.stack(norms))
    return total_norm
//  [python]
def multi_tensor_total_norm(grads, chunk_size=2048 * 32) -> torch.Tensor:
    per_device_grads = {}
    norms = []
    for grad in grads:
        device = grad.device
        cur_device_grads = per_device_grads.get(device)
        if cur_device_grads is None:
            cur_device_grads = []
            per_device_grads[device] = cur_device_grads
        cur_device_grads.append(grad)
    for device in per_device_grads.keys():
        cur_device_grads = per_device_grads[device]
        if device.type == "cuda":
            # TODO(msb) return has_inf
            has_inf = torch.zeros((1, 1), dtype=torch.int, device=device)
            with torch.cuda.device(device):
                norm = multi_tensor_l2norm(
                    chunk_size, has_inf, [cur_device_grads], False
                )
            norms.append(norm[0].to(torch.cuda.current_device()))
        else:
            norms += [torch.linalg.norm(g, ord=2, dtype=torch.float32) for g in cur_device_grads]
    total_norm = torch.linalg.norm(torch.stack(norms))
    return total_norm
//  [python]
error_array = []
eigvalues = []
nepoch_u0 = int(nepoch_u0)
start = time.time()
for epoch in range(nepoch_u0):
    if epoch ==1000:
        end = time.time()
        consume_time = end-start
        print('time is %f' % consume_time)
    if epoch%100 == 0:
        dom_koch_n = koch_points.get_koch_points(10000) # 获得n_test个koch的随机分布点
        dom_koch_n1 = dom_koch_n[(dom_koch_n[:,0]**2+dom_koch_n[:,1]**2)<r0**2]#定义内部的dom点，即是r<r0的点
        dom_koch_n2 = dom_koch_n[(dom_koch_n[:,0]**2+dom_koch_n[:,1]**2)>=r0**2]#定义内部的dom点，即是r<r0的点
        dom_koch_t1= torch.tensor(dom_koch_n1,  requires_grad=True, device = 'cuda') # 将numpy数据转化为tensor
        dom_koch_t2= torch.tensor(dom_koch_n2,  requires_grad=True, device = 'cuda') # 将numpy数据转化为tensor
        f1 = -16 * torch.norm(dom_koch_t1, dim = 1, keepdim=True).data**2 # 定义体力 
        f2 = -16 * torch.norm(dom_koch_t2, dim = 1, keepdim=True).data**2 # 定义体力
    def closure():  

        # 构造可能位移场
        u_pred1 = pred(dom_koch_t1)
        u_pred2 = pred(dom_koch_t2)
        
        du1dxy = grad(u_pred1, dom_koch_t1, torch.ones(dom_koch_t1.size()[0], 1).cuda(), create_graph=True)[0]
        du1dx = du1dxy[:, 0].unsqueeze(1)
        du1dy = du1dxy[:, 1].unsqueeze(1)

        du2dxy = grad(u_pred2, dom_koch_t2, torch.ones(dom_koch_t2.size()[0], 1).cuda(), create_graph=True)[0]
        du2dx = du2dxy[:, 0].unsqueeze(1)
        du2dy = du2dxy[:, 1].unsqueeze(1)

        J1 = ((0.5 * a1 * torch.sum(du1dx**2 + du1dy**2)) - torch.sum(f1*u_pred1)) * (r0**2*np.pi)/len(dom_koch_t1) 
        
        J2 = ((0.5 * a2 * torch.sum(du2dx**2 + du2dy**2)) - torch.sum(f2*u_pred2)) * (10*np.sqrt(3)/3-r0**2*np.pi)/len(dom_koch_t2) 

        J = J1 + J2

        loss = J 
        error_t = evaluate()
        optim_h.zero_grad()
        loss.backward()
        loss_array.append(loss.data.cpu())
        error_array.append(error_t.data.cpu())

        if epoch%10==0:
            print(' epoch : %i, the loss : %f ,  error : %f' % (epoch, loss.data, error_t.data))
        if epoch%500==0:
            x_ntk = torch.cat((dom_koch_t1[0:50], dom_koch_t2[0:50])).data
            pred_ntk = pred(x_ntk)
            eigvalue = NTK(pred_ntk)
            eigvalues.append(eigvalue)
            print('the NTK is done')
        return loss
    optim_h.step(closure)
    scheduler.step()
//  [python]
error_array = []
eigvalues = []
nepoch_u0 = int(nepoch_u0)
start = time.time()
for epoch in range(nepoch_u0):
    if epoch ==1000:
        end = time.time()
        consume_time = end-start
        print('time is %f' % consume_time)
    if epoch%100 == 0:
        dom_koch_n = koch_points.get_koch_points(10000) # 获得n_test个koch的随机分布点
        dom_koch_n1 = dom_koch_n[(dom_koch_n[:,0]**2+dom_koch_n[:,1]**2)<r0**2]#定义内部的dom点，即是r<r0的点
        dom_koch_n2 = dom_koch_n[(dom_koch_n[:,0]**2+dom_koch_n[:,1]**2)>=r0**2]#定义内部的dom点，即是r<r0的点
        dom_koch_t1= torch.tensor(dom_koch_n1,  requires_grad=True, device = 'cuda') # 将numpy数据转化为tensor
        dom_koch_t2= torch.tensor(dom_koch_n2,  requires_grad=True, device = 'cuda') # 将numpy数据转化为tensor
        f1 = -16 * torch.linalg.norm(dom_koch_t1, dim = 1, keepdim=True).data**2 # 定义体力 
        f2 = -16 * torch.linalg.norm(dom_koch_t2, dim = 1, keepdim=True).data**2 # 定义体力
    def closure():  

        # 构造可能位移场
        u_pred1 = pred(dom_koch_t1)
        u_pred2 = pred(dom_koch_t2)
        
        du1dxy = grad(u_pred1, dom_koch_t1, torch.ones(dom_koch_t1.size()[0], 1).cuda(), create_graph=True)[0]
        du1dx = du1dxy[:, 0].unsqueeze(1)
        du1dy = du1dxy[:, 1].unsqueeze(1)

        du2dxy = grad(u_pred2, dom_koch_t2, torch.ones(dom_koch_t2.size()[0], 1).cuda(), create_graph=True)[0]
        du2dx = du2dxy[:, 0].unsqueeze(1)
        du2dy = du2dxy[:, 1].unsqueeze(1)

        J1 = ((0.5 * a1 * torch.sum(du1dx**2 + du1dy**2)) - torch.sum(f1*u_pred1)) * (r0**2*np.pi)/len(dom_koch_t1) 
        
        J2 = ((0.5 * a2 * torch.sum(du2dx**2 + du2dy**2)) - torch.sum(f2*u_pred2)) * (10*np.sqrt(3)/3-r0**2*np.pi)/len(dom_koch_t2) 

        J = J1 + J2

        loss = J 
        error_t = evaluate()
        optim_h.zero_grad()
        loss.backward()
        loss_array.append(loss.data.cpu())
        error_array.append(error_t.data.cpu())

        if epoch%10==0:
            print(' epoch : %i, the loss : %f ,  error : %f' % (epoch, loss.data, error_t.data))
        if epoch%500==0:
            x_ntk = torch.cat((dom_koch_t1[0:50], dom_koch_t2[0:50])).data
            pred_ntk = pred(x_ntk)
            eigvalue = NTK(pred_ntk)
            eigvalues.append(eigvalue)
            print('the NTK is done')
        return loss
    optim_h.step(closure)
    scheduler.step()
//  [python]
def plan(self, start, goal, th_init, im, sdf, optim_params = {'method': 'gauss_newton', 'plan_time':np.inf, 'max_iters':np.inf, 'tol_err':1e-2, 'tol_delta':1e-3, 'reg':0.0}):
    assert th_init.shape[0] == self.num_traj_states and th_init.shape[1] == self.state_dim
    self.obs_factor.set_im_sdf(im, sdf)
    self.start_prior.set_mean(start)
    self.goal_prior.set_mean(goal)
    start_t = time.time()
    th_curr = th_init
    # print th_init.grad
    j = 0 #Current iteration counter
    err_per_iter = []
    err_init = self.error(th_init)
    err_old = err_init
    lam = 1e-4 #Used for LM method

    while True:
      print ("Current iteration, %d"%j)
      #Create linear system
      A, b, K = self.create_Ab_linear(th_curr)
      err_per_iter.append(err_old)
      if optim_params['method'] == 'gauss_newton':
        #get parameter update
        delta = optim_params['reg'] 
        dtheta = self.solve_linear_system(A, b, K, delta=delta, trust_region=False, method='chol')
        th_new = th_curr + dtheta
        err_new = self.error(th_new)
        err_delta = err_new - err_old
        th_curr = th_new
        err_old = err_new
        print('dtheta = %f, err_delta = %f'%(torch.norm(dtheta), err_delta)) 

      elif optim_params['method'] == 'lm':
        dtheta = self.solve_linear_system(A, b, K, delta=lam, trust_region=True, method='chol')
        th_new = th_curr + dtheta
        err_new = self.error(th_new)
        err_delta = err_new - err_old
        if err_delta < 0:
          th_curr = th_new
          lam = lam/10.0
          err_old = err_new
        else:
          lam = lam*10.0
        print('dtheta = %f, err_delta = %f, lambda= %f'%(torch.norm(dtheta), err_delta, lam)) 

      j = j + 1
      #Check convergence
      if check_convergence(dtheta, j, err_delta, optim_params['tol_err'], optim_params['tol_delta'], optim_params['max_iters']):
        break
      if time.time() - start_t > optim_params['plan_time']:
        print('Plan time over')
        break

    time_taken = time.time() - start_t
    return th_curr, err_init, err_new, err_per_iter, j, time_taken
//  [python]
def plan(self, start, goal, th_init, im, sdf, optim_params = {'method': 'gauss_newton', 'plan_time':np.inf, 'max_iters':np.inf, 'tol_err':1e-2, 'tol_delta':1e-3, 'reg':0.0}):
    assert th_init.shape[0] == self.num_traj_states and th_init.shape[1] == self.state_dim
    self.obs_factor.set_im_sdf(im, sdf)
    self.start_prior.set_mean(start)
    self.goal_prior.set_mean(goal)
    start_t = time.time()
    th_curr = th_init
    # print th_init.grad
    j = 0 #Current iteration counter
    err_per_iter = []
    err_init = self.error(th_init)
    err_old = err_init
    lam = 1e-4 #Used for LM method

    while True:
      print ("Current iteration, %d"%j)
      #Create linear system
      A, b, K = self.create_Ab_linear(th_curr)
      err_per_iter.append(err_old)
      if optim_params['method'] == 'gauss_newton':
        #get parameter update
        delta = optim_params['reg'] 
        dtheta = self.solve_linear_system(A, b, K, delta=delta, trust_region=False, method='chol')
        th_new = th_curr + dtheta
        err_new = self.error(th_new)
        err_delta = err_new - err_old
        th_curr = th_new
        err_old = err_new
        print('dtheta = %f, err_delta = %f'%(torch.linalg.norm(dtheta), err_delta)) 

      elif optim_params['method'] == 'lm':
        dtheta = self.solve_linear_system(A, b, K, delta=lam, trust_region=True, method='chol')
        th_new = th_curr + dtheta
        err_new = self.error(th_new)
        err_delta = err_new - err_old
        if err_delta < 0:
          th_curr = th_new
          lam = lam/10.0
          err_old = err_new
        else:
          lam = lam*10.0
        print('dtheta = %f, err_delta = %f, lambda= %f'%(torch.linalg.norm(dtheta), err_delta, lam)) 

      j = j + 1
      #Check convergence
      if check_convergence(dtheta, j, err_delta, optim_params['tol_err'], optim_params['tol_delta'], optim_params['max_iters']):
        break
      if time.time() - start_t > optim_params['plan_time']:
        print('Plan time over')
        break

    time_taken = time.time() - start_t
    return th_curr, err_init, err_new, err_per_iter, j, time_taken