jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def initializer(self, δ):
        Δ = torch.randn(self.X.size(), dtype=torch.double)
        Δ = Δ / torch.linalg.norm(Δ)
        return (self.X + δ * Δ).requires_grad_(True)
//  [python]
def loss(self):
        return lambda x: torch.norm(x - self.proximal_gradient(x))
//  [python]
def loss(self):
        return lambda x: torch.linalg.norm(x - self.proximal_gradient(x))
//  [python]
def norm_grad(self):
        return lambda x: torch.norm(self.loss_grad()(x))
//  [python]
def norm_grad(self):
        return lambda x: torch.linalg.norm(self.loss_grad()(x))
//  [python]
def loss(self):
        return (
            lambda x: torch.mean(torch.log(1 + torch.exp(-self.y * (self.A @ x))))
            + self.l2_penalty * torch.norm(x) ** 2
        )
//  [python]
def loss(self):
        return (
            lambda x: torch.mean(torch.log(1 + torch.exp(-self.y * (self.A @ x))))
            + self.l2_penalty * torch.linalg.norm(x) ** 2
        )
//  [python]
def loss(self):
        def f(z):
            W = z[0 : self.W.numel()].view(self.W.size())
            X = z[self.W.numel() :].view(self.X.size())
            return (1.0 / self.y.size(0)) * torch.norm(
                self.y - torch.sum(self.L.mm(W) * self.R.mm(X), dim=1), 1
            )

        return f
//  [python]
def loss(self):
        def f(z):
            W = z[0 : self.W.numel()].view(self.W.size())
            X = z[self.W.numel() :].view(self.X.size())
            return (1.0 / self.y.size(0)) * torch.linalg.norm(
                self.y - torch.sum(self.L.mm(W) * self.R.mm(X), dim=1), 1
            )

        return f
//  [python]
def _homography_joint_svd(self,
        top_corners: torch.Tensor, # in [-1, 1]
        bottom_corners: torch.Tensor, # in [-1, 1]
        floor_z: float=-1.6,
        ceil_z: float=1.6,
    ):
        b, N, _ = top_corners.size()
        floor_u = bottom_corners[:, :, 0] * np.pi
        floor_v = bottom_corners[:, :, 1] * (-0.5 * np.pi)
        floor_c = floor_z / torch.tan(floor_v)
        floor_x = floor_c * torch.sin(floor_u)
        floor_y = -floor_c * torch.cos(floor_u)
        floor_xy = torch.stack([floor_x, floor_y], dim=-1)
        floor_scale = self._get_scale_all(floor_xy)
        floor_scale = floor_scale / 2.0        
        floor_ceil_c = torch.norm(floor_xy, p=2, dim=-1)
        floor_ceil_v = top_corners[:, :, 1] * (-0.5 * np.pi)
        floor_ceil_z = (floor_ceil_c * torch.tan(floor_ceil_v)).mean(dim=1, keepdim=True)
        floor_ceil_z = floor_ceil_z.unsqueeze(1).expand(b, 4, 1).contiguous()
        ceil_u_t = top_corners[:, :, 0] * np.pi
        ceil_v_t = top_corners[:, :, 1] * (-0.5 * np.pi)
        ceil_c = ceil_z / torch.tan(ceil_v_t)
        ceil_x = ceil_c * torch.sin(ceil_u_t)
        ceil_y = -ceil_c * torch.cos(ceil_u_t)
        ceil_xy = torch.stack([ceil_x, ceil_y], dim=-1)
        ceil_floor_c = torch.norm(ceil_xy, p=2, dim=-1)
        ceil_v_b = bottom_corners[:, :, 1] * (-0.5 * np.pi)
        ceil_floor_z = (ceil_floor_c * torch.tan(ceil_v_b)).mean(dim=1, keepdim=True)
        fix_ceil = -ceil_z / ceil_floor_z
        ceil_z_fix = ceil_z * fix_ceil
        ceil_z_fix = ceil_z_fix.unsqueeze(1).expand(b, 4, 1).contiguous()
        ceil_floor_fixed_c = ceil_z_fix.squeeze(-1) / torch.tan(ceil_v_t)
        ceil_x = ceil_floor_fixed_c * torch.sin(ceil_u_t)
        ceil_y = -ceil_floor_fixed_c * torch.cos(ceil_u_t)
        ceil_xy = torch.stack([ceil_x, ceil_y], dim=-1)
        ceil_scale = self._get_scale_all(ceil_xy)
        ceil_scale = ceil_scale / 2.0
        joint_xy = 0.5 * (floor_xy + ceil_xy)
        joint_scale = 0.5 * (floor_scale + ceil_scale)        
        joint_centroid = joint_xy.mean(dim=1)
        joint_xy = joint_xy - joint_centroid.unsqueeze(1)
        inds = torch.sort(torch.atan2(joint_xy[..., 0], joint_xy[..., 1] + 1e-12))[1]
        axes = self.cuboid_axes[:, inds.squeeze(), :]
        homography = kornia.get_perspective_transform(joint_xy, axes)    
        homogeneous = torch.cat([joint_xy, torch.ones_like(joint_xy[..., -1:])], dim=2)
        xformed = (homography @ homogeneous.transpose(1, 2)).transpose(1, 2)
        xformed = xformed[:, :, :2] / xformed[:, :, 2].unsqueeze(-1)
        rect_joint_xy = xformed * joint_scale.unsqueeze(1) + joint_centroid.unsqueeze(1)
        original_xy = joint_xy + joint_centroid.unsqueeze(1)
        R, t, s = self._svd(rect_joint_xy, original_xy[:, inds.squeeze(), :])
        rect_joint_xy = self._transform_points(rect_joint_xy, R, t, s)
        bottom_points = torch.cat([rect_joint_xy, floor_z * torch.ones_like(floor_c.unsqueeze(-1))], dim=-1)
        top_points = torch.cat([rect_joint_xy, ceil_z_fix], dim=-1)
        return top_points, bottom_points
//  [python]
def _homography_joint_svd(self,
        top_corners: torch.Tensor, # in [-1, 1]
        bottom_corners: torch.Tensor, # in [-1, 1]
        floor_z: float=-1.6,
        ceil_z: float=1.6,
    ):
        b, N, _ = top_corners.size()
        floor_u = bottom_corners[:, :, 0] * np.pi
        floor_v = bottom_corners[:, :, 1] * (-0.5 * np.pi)
        floor_c = floor_z / torch.tan(floor_v)
        floor_x = floor_c * torch.sin(floor_u)
        floor_y = -floor_c * torch.cos(floor_u)
        floor_xy = torch.stack([floor_x, floor_y], dim=-1)
        floor_scale = self._get_scale_all(floor_xy)
        floor_scale = floor_scale / 2.0        
        floor_ceil_c = torch.linalg.norm(floor_xy, ord=2, dim=-1)
        floor_ceil_v = top_corners[:, :, 1] * (-0.5 * np.pi)
        floor_ceil_z = (floor_ceil_c * torch.tan(floor_ceil_v)).mean(dim=1, keepdim=True)
        floor_ceil_z = floor_ceil_z.unsqueeze(1).expand(b, 4, 1).contiguous()
        ceil_u_t = top_corners[:, :, 0] * np.pi
        ceil_v_t = top_corners[:, :, 1] * (-0.5 * np.pi)
        ceil_c = ceil_z / torch.tan(ceil_v_t)
        ceil_x = ceil_c * torch.sin(ceil_u_t)
        ceil_y = -ceil_c * torch.cos(ceil_u_t)
        ceil_xy = torch.stack([ceil_x, ceil_y], dim=-1)
        ceil_floor_c = torch.linalg.norm(ceil_xy, ord=2, dim=-1)
        ceil_v_b = bottom_corners[:, :, 1] * (-0.5 * np.pi)
        ceil_floor_z = (ceil_floor_c * torch.tan(ceil_v_b)).mean(dim=1, keepdim=True)
        fix_ceil = -ceil_z / ceil_floor_z
        ceil_z_fix = ceil_z * fix_ceil
        ceil_z_fix = ceil_z_fix.unsqueeze(1).expand(b, 4, 1).contiguous()
        ceil_floor_fixed_c = ceil_z_fix.squeeze(-1) / torch.tan(ceil_v_t)
        ceil_x = ceil_floor_fixed_c * torch.sin(ceil_u_t)
        ceil_y = -ceil_floor_fixed_c * torch.cos(ceil_u_t)
        ceil_xy = torch.stack([ceil_x, ceil_y], dim=-1)
        ceil_scale = self._get_scale_all(ceil_xy)
        ceil_scale = ceil_scale / 2.0
        joint_xy = 0.5 * (floor_xy + ceil_xy)
        joint_scale = 0.5 * (floor_scale + ceil_scale)        
        joint_centroid = joint_xy.mean(dim=1)
        joint_xy = joint_xy - joint_centroid.unsqueeze(1)
        inds = torch.sort(torch.atan2(joint_xy[..., 0], joint_xy[..., 1] + 1e-12))[1]
        axes = self.cuboid_axes[:, inds.squeeze(), :]
        homography = kornia.get_perspective_transform(joint_xy, axes)    
        homogeneous = torch.cat([joint_xy, torch.ones_like(joint_xy[..., -1:])], dim=2)
        xformed = (homography @ homogeneous.transpose(1, 2)).transpose(1, 2)
        xformed = xformed[:, :, :2] / xformed[:, :, 2].unsqueeze(-1)
        rect_joint_xy = xformed * joint_scale.unsqueeze(1) + joint_centroid.unsqueeze(1)
        original_xy = joint_xy + joint_centroid.unsqueeze(1)
        R, t, s = self._svd(rect_joint_xy, original_xy[:, inds.squeeze(), :])
        rect_joint_xy = self._transform_points(rect_joint_xy, R, t, s)
        bottom_points = torch.cat([rect_joint_xy, floor_z * torch.ones_like(floor_c.unsqueeze(-1))], dim=-1)
        top_points = torch.cat([rect_joint_xy, ceil_z_fix], dim=-1)
        return top_points, bottom_points
//  [python]
def test_sqrt_sum_with_cpu_random_data(test_case):
        device = cpu_device()
        x = random_tensor(ndim=4, dim1=3, dim2=4, dim3=5, requires_grad=False).to(
            device
        )
        y = torch.norm(x)
        return y
//  [python]
def test_sqrt_sum_with_cpu_random_data(test_case):
        device = cpu_device()
        x = random_tensor(ndim=4, dim1=3, dim2=4, dim3=5, requires_grad=False).to(
            device
        )
        y = torch.linalg.norm(x)
        return y
//  [python]
def test_sqrt_sum_with_cuda_random_data(test_case):
        device = gpu_device()
        x = random_tensor(ndim=4, dim1=10, dim2=10, dim3=10, requires_grad=False).to(
            device
        )
        y = x.norm()
        return y
//  [python]
def test_sqrt_sum_with_cuda_random_data(test_case):
        device = gpu_device()
        x = random_tensor(ndim=4, dim1=10, dim2=10, dim3=10, requires_grad=False).to(
            device
        )
        y = torch.linalg.norm(x)
        return y
//  [python]
def test_scalar_print_random_data(test_case):
        device = random_device()
        x = random_tensor(ndim=4, dim1=3, dim2=4, dim3=5, requires_grad=False).to(
            device
        )
        y = torch.norm(x)
        return y
//  [python]
def test_scalar_print_random_data(test_case):
        device = random_device()
        x = random_tensor(ndim=4, dim1=3, dim2=4, dim3=5, requires_grad=False).to(
            device
        )
        y = torch.linalg.norm(x)
        return y
//  [python]
def generate_rays(self, pixels_xys, intrinsics, c2ws):
        """

        Parameters
        ----------
        pixels_xys : [N_views, N_num, 2]
        intrinsics : [N_views, 3, 3]
        c2ws : [N_views, 4, 4]

        Returns
        -------

        """
        intrinsics_inv = torch.inverse(intrinsics)
        p = torch.cat([pixels_xys, torch.ones_like(pixels_xys[:, :, :1])], dim=-1)  # N_views, n, 3
        p = torch.matmul(intrinsics_inv[:, None, :3, :3], p[:, :, :, None]).squeeze()  # N_views, n, 3
        rays_v = p / torch.norm(p, p=2, dim=-1, keepdim=True)  # N_views, n, 3
        rays_v = torch.matmul(c2ws[:, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # N_views, n, 3
        rays_o = c2ws[:, None, :3, 3].expand(rays_v.shape)  # N_views, n, 3
        return rays_o, rays_v
//  [python]
def generate_rays(self, pixels_xys, intrinsics, c2ws):
        """

        Parameters
        ----------
        pixels_xys : [N_views, N_num, 2]
        intrinsics : [N_views, 3, 3]
        c2ws : [N_views, 4, 4]

        Returns
        -------

        """
        intrinsics_inv = torch.inverse(intrinsics)
        p = torch.cat([pixels_xys, torch.ones_like(pixels_xys[:, :, :1])], dim=-1)  # N_views, n, 3
        p = torch.matmul(intrinsics_inv[:, None, :3, :3], p[:, :, :, None]).squeeze()  # N_views, n, 3
        rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # N_views, n, 3
        rays_v = torch.matmul(c2ws[:, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # N_views, n, 3
        rays_o = c2ws[:, None, :3, 3].expand(rays_v.shape)  # N_views, n, 3
        return rays_o, rays_v
//  [python]
def conditioning(self, x_prev, x_t, x_0_hat, measurement, **kwargs):
        norm = 0
        for _ in range(self.num_sampling):
            # TODO: use noiser?
            x_0_hat_noise = x_0_hat + 0.05 * torch.rand_like(x_0_hat)
            difference = measurement - self.operator.forward(x_0_hat_noise)
            norm += torch.norm(difference) / self.num_sampling
        
        norm_grad = torch.autograd.grad(outputs=norm, inputs=x_prev)[0]
        x_t -= norm_grad * self.scale
        return x_t, norm
//  [python]
def conditioning(self, x_prev, x_t, x_0_hat, measurement, **kwargs):
        norm = 0
        for _ in range(self.num_sampling):
            # TODO: use noiser?
            x_0_hat_noise = x_0_hat + 0.05 * torch.rand_like(x_0_hat)
            difference = measurement - self.operator.forward(x_0_hat_noise)
            norm += torch.linalg.norm(difference) / self.num_sampling
        
        norm_grad = torch.autograd.grad(outputs=norm, inputs=x_prev)[0]
        x_t -= norm_grad * self.scale
        return x_t, norm
//  [python]
def grad_and_value(self, x_prev, x_0_hat, measurement, **kwargs):
        if self.noiser.__name__ == 'gaussian':
            difference = measurement - self.operator.forward(x_0_hat, **kwargs)
            norm = torch.norm(difference)
            norm_grad = torch.autograd.grad(outputs=norm, inputs=x_prev)[0]
        
        elif self.noiser.__name__ == 'poisson':
            Ax = self.operator.forward(x_0_hat, **kwargs)
            difference = measurement-Ax
            norm = torch.norm(difference) / measurement.abs()
            norm = norm.mean()
            norm_grad = torch.autograd.grad(outputs=norm, inputs=x_prev)[0]

        else:
            raise NotImplementedError
             
        return norm_grad, norm
//  [python]
def grad_and_value(self, x_prev, x_0_hat, measurement, **kwargs):
        if self.noiser.__name__ == 'gaussian':
            difference = measurement - self.operator.forward(x_0_hat, **kwargs)
            norm = torch.linalg.norm(difference)
            norm_grad = torch.autograd.grad(outputs=norm, inputs=x_prev)[0]
        
        elif self.noiser.__name__ == 'poisson':
            Ax = self.operator.forward(x_0_hat, **kwargs)
            difference = measurement-Ax
            norm = torch.linalg.norm(difference) / measurement.abs()
            norm = norm.mean()
            norm_grad = torch.autograd.grad(outputs=norm, inputs=x_prev)[0]

        else:
            raise NotImplementedError
             
        return norm_grad, norm
//  [python]
def regularize(point):
    """Regularize a point to be in accordance with convention.
    In 3D, regularize the norm of the rotation vector,
    to be between 0 and pi, following the axis-angle
    representation's convention.
    If the angle is between pi and 2pi,
    the function computes its complementary in 2pi and
    inverts the direction of the rotation axis.
    Parameters

    # Adapted from geomstats
    # https://github.com/geomstats/geomstats/blob/master/geomstats/geometry/special_orthogonal.py#L884
    ----------
    point : array-like, shape=[...,3]
        Point.
    Returns
    -------
    regularized_point : array-like, shape=[..., 3]
        Regularized point.
    """
    theta = torch.norm(point, dim=-1)
    k = torch.floor(theta / 2.0 / torch.pi)

    # angle in [0;2pi)
    angle = theta - 2 * k * torch.pi

    # this avoids dividing by 0
    theta_eps = torch.where(torch.isclose(theta, torch.tensor(0.0)), 1.0, theta)

    # angle in [0, pi]
    normalized_angle = torch.where(angle <= torch.pi, angle, 2 * torch.pi - angle)
    norm_ratio = torch.where(torch.isclose(theta, torch.tensor(0.0)), 1.0, normalized_angle / theta_eps)

    # reverse sign if angle was greater than pi
    norm_ratio = torch.where(angle > torch.pi, -norm_ratio, norm_ratio)
    return torch.einsum("...,...i->...i", norm_ratio, point)
//  [python]
def regularize(point):
    """Regularize a point to be in accordance with convention.
    In 3D, regularize the norm of the rotation vector,
    to be between 0 and pi, following the axis-angle
    representation's convention.
    If the angle is between pi and 2pi,
    the function computes its complementary in 2pi and
    inverts the direction of the rotation axis.
    Parameters

    # Adapted from geomstats
    # https://github.com/geomstats/geomstats/blob/master/geomstats/geometry/special_orthogonal.py#L884
    ----------
    point : array-like, shape=[...,3]
        Point.
    Returns
    -------
    regularized_point : array-like, shape=[..., 3]
        Regularized point.
    """
    theta = torch.linalg.norm(point, axis=-1)
    k = torch.floor(theta / 2.0 / torch.pi)

    # angle in [0;2pi)
    angle = theta - 2 * k * torch.pi

    # this avoids dividing by 0
    theta_eps = torch.where(torch.isclose(theta, torch.tensor(0.0)), 1.0, theta)

    # angle in [0, pi]
    normalized_angle = torch.where(angle <= torch.pi, angle, 2 * torch.pi - angle)
    norm_ratio = torch.where(torch.isclose(theta, torch.tensor(0.0)), 1.0, normalized_angle / theta_eps)

    # reverse sign if angle was greater than pi
    norm_ratio = torch.where(angle > torch.pi, -norm_ratio, norm_ratio)
    return torch.einsum("...,...i->...i", norm_ratio, point)
//  [python]
def dot_accuracy(ns, net) -> float:
    """Calculate accuracy in the sense of normalized dot product mapped from 0 to 1.
    
    For two vectors `x` and `y`, we consider the following measure of normalized dot
    product:
        0.5 * (1 + dot(x, y) / (norm(x) * norm(y))) .

    :param ns: namespace with output from `TrainingBatch.feed()`; should have fields
        `ns.y` (the ground truth) and `ns.y_pred` (the prediction)
    :param net: network making the prediction; this is unused here
    :return: an indication of the normalized dot product (uncentered correlation)
        averaged over samples; this will range from 0 to 1
    """
    y = ns.y
    y_pred = ns.y_pred

    norm_pred = torch.norm(y_pred, dim=1)
    norm = torch.norm(y, dim=1)

    dot = torch.sum(y_pred * y, dim=1)

    accuracy = 0.5 * (1 + dot / (norm_pred * norm))
    return torch.mean(accuracy).item()
//  [python]
def dot_accuracy(ns, net) -> float:
    """Calculate accuracy in the sense of normalized dot product mapped from 0 to 1.
    
    For two vectors `x` and `y`, we consider the following measure of normalized dot
    product:
        0.5 * (1 + dot(x, y) / (norm(x) * norm(y))) .

    :param ns: namespace with output from `TrainingBatch.feed()`; should have fields
        `ns.y` (the ground truth) and `ns.y_pred` (the prediction)
    :param net: network making the prediction; this is unused here
    :return: an indication of the normalized dot product (uncentered correlation)
        averaged over samples; this will range from 0 to 1
    """
    y = ns.y
    y_pred = ns.y_pred

    norm_pred = torch.linalg.norm(y_pred, dim=1)
    norm = torch.linalg.norm(y, dim=1)

    dot = torch.sum(y_pred * y, dim=1)

    accuracy = 0.5 * (1 + dot / (norm_pred * norm))
    return torch.mean(accuracy).item()
//  [python]
def map_cartesian_to_polar(self, z):
        """
        translates bxd sized tensor interpreted as cartesian coordinates to b x (d+1) polar coordinates.
        The first entry per sample is the radius, the other are angles from the individual axis.
        Technically, the polar coordinates are over-defined, but that way all the signs are recoverable.
        """
        from einops import repeat

        # compute radius
        # r = sqrt(sum_i(z_i**2))
        r_norm = torch.norm(z, p=2, dim=1)
        # apply trigonometry:
        # z_j = r cos(theta_j) <=> theta_j = arccos(z_j/r)
        r_norm_rep = repeat(r_norm, "b -> b d", d=z.shape[1])
        r = torch.acos(z / r_norm_rep)
        # cat the radius and angles
        r = torch.cat((r_norm.unsqueeze(dim=1), r), dim=1)
        return r
//  [python]
def map_cartesian_to_polar(self, z):
        """
        translates bxd sized tensor interpreted as cartesian coordinates to b x (d+1) polar coordinates.
        The first entry per sample is the radius, the other are angles from the individual axis.
        Technically, the polar coordinates are over-defined, but that way all the signs are recoverable.
        """
        from einops import repeat

        # compute radius
        # r = sqrt(sum_i(z_i**2))
        r_norm = torch.linalg.norm(z, ord=2, dim=1)
        # apply trigonometry:
        # z_j = r cos(theta_j) <=> theta_j = arccos(z_j/r)
        r_norm_rep = repeat(r_norm, "b -> b d", d=z.shape[1])
        r = torch.acos(z / r_norm_rep)
        # cat the radius and angles
        r = torch.cat((r_norm.unsqueeze(dim=1), r), dim=1)
        return r
//  [python]
def forward(self, x):
        self.model.eval()

        with torch.no_grad():
            logits = self.model(x)
            features = self.model.get_features()
            norm = torch.norm(features, dim=-1, keepdim=True)
            features = features / norm

            # Forward with feature normalization
            logits_w_norm_features = self.w(features)
        return logits, features, logits_w_norm_features
//  [python]
def forward(self, x):
        self.model.eval()

        with torch.no_grad():
            logits = self.model(x)
            features = self.model.get_features()
            norm = torch.linalg.norm(features, dim=-1, keepdim=True)
            features = features / norm

            # Forward with feature normalization
            logits_w_norm_features = self.w(features)
        return logits, features, logits_w_norm_features
//  [python]
def add_edge_features(self) -> List[dgl.DGLGraph]:
        """
        Add relative displacement and displacement norm as edge features for each graph
        in the list of graphs. The calculations are done using the 'pos' attribute in the
        node data of each graph. The resulting edge features are stored in the 'x' attribute
        in the edge data of each graph.

        This method will modify the list of graphs in-place.

        Returns
        -------
        List[dgl.DGLGraph]
            The list of graphs with updated edge features.
        """
        if not hasattr(self, "graphs") or not self.graphs:
            raise ValueError("The list 'graphs' is empty.")

        for graph in self.graphs:
            pos = graph.ndata.get("pos")
            if pos is None:
                raise ValueError(
                    "'pos' does not exist in the node data of one or more graphs."
                )

            row, col = graph.edges()
            row = row.long()
            col = col.long()

            disp = pos[row] - pos[col]
            disp_norm = torch.norm(disp, dim=-1, keepdim=True)
            graph.edata["x"] = torch.cat((disp, disp_norm), dim=-1)

        return self.graphs
//  [python]
def add_edge_features(self) -> List[dgl.DGLGraph]:
        """
        Add relative displacement and displacement norm as edge features for each graph
        in the list of graphs. The calculations are done using the 'pos' attribute in the
        node data of each graph. The resulting edge features are stored in the 'x' attribute
        in the edge data of each graph.

        This method will modify the list of graphs in-place.

        Returns
        -------
        List[dgl.DGLGraph]
            The list of graphs with updated edge features.
        """
        if not hasattr(self, "graphs") or not self.graphs:
            raise ValueError("The list 'graphs' is empty.")

        for graph in self.graphs:
            pos = graph.ndata.get("pos")
            if pos is None:
                raise ValueError(
                    "'pos' does not exist in the node data of one or more graphs."
                )

            row, col = graph.edges()
            row = row.long()
            col = col.long()

            disp = pos[row] - pos[col]
            disp_norm = torch.linalg.norm(disp, dim=-1, keepdim=True)
            graph.edata["x"] = torch.cat((disp, disp_norm), dim=-1)

        return self.graphs
//  [python]
def test_relearn_structure(self, wordvec_pretrain_file):
        """
        Test that starting a trainer with --relearn_structure copies the old model
        """

        with tempfile.TemporaryDirectory(dir=TEST_WORKING_DIR) as tmpdirname:
            set_random_seed(1000)
            args = ['--pattn_num_layers', '0', '--lattn_d_proj', '0', '--hidden_size', '20', '--delta_embedding_dim', '10']
            tr = build_trainer(wordvec_pretrain_file, *args)

            # attempt saving
            filename = os.path.join(tmpdirname, "parser.pt")
            tr.save(filename)

            set_random_seed(1001)
            args = ['--pattn_num_layers', '1', '--lattn_d_proj', '0', '--hidden_size', '20', '--delta_embedding_dim', '10', '--relearn_structure', '--load_name', filename]
            tr2 = build_trainer(wordvec_pretrain_file, *args)

            assert torch.allclose(tr.model.delta_embedding.weight, tr2.model.delta_embedding.weight)
            assert torch.allclose(tr.model.output_layers[0].weight, tr2.model.output_layers[0].weight)
            # the norms will be the same, as the non-zero values are all the same
            assert torch.allclose(torch.norm(tr.model.word_lstm.weight_ih_l0), torch.norm(tr2.model.word_lstm.weight_ih_l0))
//  [python]
def test_relearn_structure(self, wordvec_pretrain_file):
        """
        Test that starting a trainer with --relearn_structure copies the old model
        """

        with tempfile.TemporaryDirectory(dir=TEST_WORKING_DIR) as tmpdirname:
            set_random_seed(1000)
            args = ['--pattn_num_layers', '0', '--lattn_d_proj', '0', '--hidden_size', '20', '--delta_embedding_dim', '10']
            tr = build_trainer(wordvec_pretrain_file, *args)

            # attempt saving
            filename = os.path.join(tmpdirname, "parser.pt")
            tr.save(filename)

            set_random_seed(1001)
            args = ['--pattn_num_layers', '1', '--lattn_d_proj', '0', '--hidden_size', '20', '--delta_embedding_dim', '10', '--relearn_structure', '--load_name', filename]
            tr2 = build_trainer(wordvec_pretrain_file, *args)

            assert torch.allclose(tr.model.delta_embedding.weight, tr2.model.delta_embedding.weight)
            assert torch.allclose(tr.model.output_layers[0].weight, tr2.model.output_layers[0].weight)
            # the norms will be the same, as the non-zero values are all the same
            assert torch.allclose(torch.linalg.norm(tr.model.word_lstm.weight_ih_l0), torch.linalg.norm(tr2.model.word_lstm.weight_ih_l0))
//  [python]
def training_step(self, batch, batch_idx):
        x, y = batch
        logits, features, logits_w_norm_features = self.forward(x)
        probs = torch.nn.functional.softmax(logits, dim=-1)
        probs_normalized = torch.nn.functional.softmax(logits_w_norm_features, dim=-1)
        loss = F.cross_entropy(logits, y)

        y_hat = probs.argmax(dim=-1)
        is_correct = y_hat == y
        output = {
            "probs": probs,
            "probs_normalized": probs_normalized,
            "features": features,
            "loss": loss,
            "is_correct": is_correct,
            "logits": logits,
            "norms": torch.norm(features, dim=-1),
        }
        return output
//  [python]
def training_step(self, batch, batch_idx):
        x, y = batch
        logits, features, logits_w_norm_features = self.forward(x)
        probs = torch.nn.functional.softmax(logits, dim=-1)
        probs_normalized = torch.nn.functional.softmax(logits_w_norm_features, dim=-1)
        loss = F.cross_entropy(logits, y)

        y_hat = probs.argmax(dim=-1)
        is_correct = y_hat == y
        output = {
            "probs": probs,
            "probs_normalized": probs_normalized,
            "features": features,
            "loss": loss,
            "is_correct": is_correct,
            "logits": logits,
            "norms": torch.linalg.norm(features, dim=-1),
        }
        return output
//  [python]
def test_step(self, batch, batch_idx):
        x, y = batch
        logits, features, logits_w_norm_features = self.forward(x)
        probs = torch.nn.functional.softmax(logits, dim=-1)
        probs_normalized = torch.nn.functional.softmax(logits_w_norm_features, dim=-1)

        y_hat = probs.argmax(dim=-1)
        is_correct = y_hat == y

        output = {
            "probs": probs,
            "probs_normalized": probs_normalized,
            "is_correct": is_correct,
            "features": features,
            "logits": logits,
            "norms": torch.norm(features, dim=-1),
        }
        return output
//  [python]
def test_step(self, batch, batch_idx):
        x, y = batch
        logits, features, logits_w_norm_features = self.forward(x)
        probs = torch.nn.functional.softmax(logits, dim=-1)
        probs_normalized = torch.nn.functional.softmax(logits_w_norm_features, dim=-1)

        y_hat = probs.argmax(dim=-1)
        is_correct = y_hat == y

        output = {
            "probs": probs,
            "probs_normalized": probs_normalized,
            "is_correct": is_correct,
            "features": features,
            "logits": logits,
            "norms": torch.linalg.norm(features, dim=-1),
        }
        return output
//  [python]
def render_img(self, img_idx, resolution):
        pixels, radii, _, cam_intr_0, cam_pose_0 = self.dataset.gen_rays_at(img_idx, resolution_level=resolution)
        if self.camera_model is not None:
            cam_intr_, cam_pose_ = self.camera_model(img_idx, cam_intr_0, cam_pose_0)
        else:
            cam_intr_ = cam_intr_0
            cam_pose_ = cam_pose_0

        p = torch.matmul(torch.linalg.inv(cam_intr_)[None, None, :3, :3], pixels[:, :, :, None]).squeeze()  # W, H, 3
        rays_d = p / torch.norm(p, p=2, dim=-1, keepdim=True)  # W, H, 3
        rays_d = torch.matmul(cam_pose_[None, None, :3, :3], rays_d[:, :, :, None]).squeeze()  # W, H, 3
        rays_o = cam_pose_[None, None, :3, 3].expand(rays_d.shape)  # W, H, 3
        
        H, W, _ = rays_o.shape
        rays_o = rays_o.transpose(0, 1).reshape(-1, 3).split(800)
        rays_d = rays_d.transpose(0, 1).reshape(-1, 3).split(800)

        radii = radii.reshape(-1, 1).split(800)

        out_rgb_fine = []
        for rays_o_batch, rays_d_batch, radii_batch in tqdm(zip(rays_o, rays_d, radii)):
            near, far = self.dataset.near_far_from_sphere(rays_o_batch, rays_d_batch)
            background_rgb = torch.ones([1, 3]) if self.use_white_bkgd else None

            render_out = self.renderer.render(rays_o_batch,
                                              rays_d_batch,
                                              radii_batch,
                                              near,
                                              far,
                                              cos_anneal_ratio=self.get_cos_anneal_ratio(),
                                              background_rgb=background_rgb, perturb_overwrite=0)

            out_rgb_fine.append(render_out['color_fine'].detach().cpu().numpy())

            del render_out

        img_fine = (np.concatenate(out_rgb_fine, axis=0).reshape([H, W, 3]) * 256).clip(0, 255).astype(np.uint8)
        img_dir = os.path.join(self.base_exp_dir, 'render')
        os.makedirs(img_dir, exist_ok=True)
        cv.imwrite(os.path.join(img_dir, f'{self.iter_step}_{img_idx}.png'), img_fine)
//  [python]
def render_img(self, img_idx, resolution):
        pixels, radii, _, cam_intr_0, cam_pose_0 = self.dataset.gen_rays_at(img_idx, resolution_level=resolution)
        if self.camera_model is not None:
            cam_intr_, cam_pose_ = self.camera_model(img_idx, cam_intr_0, cam_pose_0)
        else:
            cam_intr_ = cam_intr_0
            cam_pose_ = cam_pose_0

        p = torch.matmul(torch.linalg.inv(cam_intr_)[None, None, :3, :3], pixels[:, :, :, None]).squeeze()  # W, H, 3
        rays_d = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # W, H, 3
        rays_d = torch.matmul(cam_pose_[None, None, :3, :3], rays_d[:, :, :, None]).squeeze()  # W, H, 3
        rays_o = cam_pose_[None, None, :3, 3].expand(rays_d.shape)  # W, H, 3
        
        H, W, _ = rays_o.shape
        rays_o = rays_o.transpose(0, 1).reshape(-1, 3).split(800)
        rays_d = rays_d.transpose(0, 1).reshape(-1, 3).split(800)

        radii = radii.reshape(-1, 1).split(800)

        out_rgb_fine = []
        for rays_o_batch, rays_d_batch, radii_batch in tqdm(zip(rays_o, rays_d, radii)):
            near, far = self.dataset.near_far_from_sphere(rays_o_batch, rays_d_batch)
            background_rgb = torch.ones([1, 3]) if self.use_white_bkgd else None

            render_out = self.renderer.render(rays_o_batch,
                                              rays_d_batch,
                                              radii_batch,
                                              near,
                                              far,
                                              cos_anneal_ratio=self.get_cos_anneal_ratio(),
                                              background_rgb=background_rgb, perturb_overwrite=0)

            out_rgb_fine.append(render_out['color_fine'].detach().cpu().numpy())

            del render_out

        img_fine = (np.concatenate(out_rgb_fine, axis=0).reshape([H, W, 3]) * 256).clip(0, 255).astype(np.uint8)
        img_dir = os.path.join(self.base_exp_dir, 'render')
        os.makedirs(img_dir, exist_ok=True)
        cv.imwrite(os.path.join(img_dir, f'{self.iter_step}_{img_idx}.png'), img_fine)
//  [python]
def _get_random_rays(num_rays, max_box, aabb, device):
    """
    generate random rays that distribute in the general direction of the box.
    The implementation is naive for now.
    :param num_rays:
    :param max_box:
    :param aabb:
    :param device:
    :return:
    """
    origins = []
    directions = []
    aabb_center = torch.tensor(
        [0.5 * (aabb[0] + aabb[3]), 0.5 * (aabb[1] + aabb[4]), 0.5 * (aabb[2] + aabb[5])],
        dtype=torch.float32,
        device=device,
    )

    for _ in range(num_rays):
        x = random.uniform(-max_box, max_box)
        y = random.uniform(-max_box, max_box)
        z = random.uniform(-max_box, max_box)

        origin = torch.tensor([x, y, z], dtype=torch.float32, device=device)
        dirs_to_cent = aabb_center - origin
        dirs_to_cent = dirs_to_cent / torch.norm(dirs_to_cent)
        phi = math.asin(dirs_to_cent[2])
        theta = torch.atan2(dirs_to_cent[1], dirs_to_cent[0])

        error_ang_phi = random.gauss(0, 0.05)
        error_ang_theta = random.gauss(0, 0.1)

        phi = phi + error_ang_phi
        theta = theta + error_ang_theta

        x_dir = math.cos(phi) * math.cos(theta)
        y_dir = math.cos(phi) * math.sin(theta)
        z_dir = math.sin(phi)

        origins.append([x, y, z])
        directions.append([x_dir, y_dir, z_dir])

    origins = torch.tensor(origins, dtype=torch.float32, device=device)
    directions = torch.tensor(directions, dtype=torch.float32, device=device)

    return origins, directions
//  [python]
def _get_random_rays(num_rays, max_box, aabb, device):
    """
    generate random rays that distribute in the general direction of the box.
    The implementation is naive for now.
    :param num_rays:
    :param max_box:
    :param aabb:
    :param device:
    :return:
    """
    origins = []
    directions = []
    aabb_center = torch.tensor(
        [0.5 * (aabb[0] + aabb[3]), 0.5 * (aabb[1] + aabb[4]), 0.5 * (aabb[2] + aabb[5])],
        dtype=torch.float32,
        device=device,
    )

    for _ in range(num_rays):
        x = random.uniform(-max_box, max_box)
        y = random.uniform(-max_box, max_box)
        z = random.uniform(-max_box, max_box)

        origin = torch.tensor([x, y, z], dtype=torch.float32, device=device)
        dirs_to_cent = aabb_center - origin
        dirs_to_cent = dirs_to_cent / torch.linalg.norm(dirs_to_cent)
        phi = math.asin(dirs_to_cent[2])
        theta = torch.atan2(dirs_to_cent[1], dirs_to_cent[0])

        error_ang_phi = random.gauss(0, 0.05)
        error_ang_theta = random.gauss(0, 0.1)

        phi = phi + error_ang_phi
        theta = theta + error_ang_theta

        x_dir = math.cos(phi) * math.cos(theta)
        y_dir = math.cos(phi) * math.sin(theta)
        z_dir = math.sin(phi)

        origins.append([x, y, z])
        directions.append([x_dir, y_dir, z_dir])

    origins = torch.tensor(origins, dtype=torch.float32, device=device)
    directions = torch.tensor(directions, dtype=torch.float32, device=device)

    return origins, directions
//  [python]
def gen_rays_pose(self, pose, resolution_level=1):
        """
        Generate rays at world space given pose.
        """
        l = resolution_level
        tx = torch.linspace(0, self.W - 1, int(self.W // l))
        ty = torch.linspace(0, self.H - 1, int(self.H // l))
        pixels_x, pixels_y = torch.meshgrid(tx, ty)
        pixels_x = pixels_x.t()
        pixels_y = pixels_y.t()
        p = torch.stack([(pixels_x - self.K[0][2]) / self.K[0][0],
                         -(pixels_y - self.K[1][2]) / self.K[1][1],
                         -torch.ones_like(pixels_x)], -1).float()
        rays_v = p / torch.norm(p, p=2, dim=-1, keepdim=True)  # W, H, 3
        rays_v = torch.sum(rays_v[..., None, :] * pose[:3, :3], -1)
        rays_o = pose[None, None, :3, 3].expand(rays_v.shape)  # W, H, 3
        return rays_o, rays_v
//  [python]
def gen_rays_pose(self, pose, resolution_level=1):
        """
        Generate rays at world space given pose.
        """
        l = resolution_level
        tx = torch.linspace(0, self.W - 1, int(self.W // l))
        ty = torch.linspace(0, self.H - 1, int(self.H // l))
        pixels_x, pixels_y = torch.meshgrid(tx, ty)
        pixels_x = pixels_x.t()
        pixels_y = pixels_y.t()
        p = torch.stack([(pixels_x - self.K[0][2]) / self.K[0][0],
                         -(pixels_y - self.K[1][2]) / self.K[1][1],
                         -torch.ones_like(pixels_x)], -1).float()
        rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # W, H, 3
        rays_v = torch.sum(rays_v[..., None, :] * pose[:3, :3], -1)
        rays_o = pose[None, None, :3, 3].expand(rays_v.shape)  # W, H, 3
        return rays_o, rays_v
//  [python]
def gen_rays_at(self, img_idx, resolution_level=1):
        """
        Generate rays at world space from one camera.
        """
        l = resolution_level
        tx = torch.linspace(0, self.W - 1, int(self.W // l))
        ty = torch.linspace(0, self.H - 1, int(self.H // l))
        pixels_x, pixels_y = torch.meshgrid(tx, ty)
        pixels_x = pixels_x.t()
        pixels_y = pixels_y.t()
        p = torch.stack([(pixels_x - self.K[0][2]) / self.K[0][0],
                         -(pixels_y - self.K[1][2]) / self.K[1][1],
                         -torch.ones_like(pixels_x)], -1).float()
        rays_v = p / torch.norm(p, p=2, dim=-1, keepdim=True)  # W, H, 3
        # rays_v = torch.matmul(self.poses[img_idx, None, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # W, H, 3
        rays_v = torch.sum(rays_v[..., None, :] * self.poses[img_idx, :3, :3], -1)
        rays_o = self.poses[img_idx, None, None, :3, 3].expand(rays_v.shape)  # W, H, 3
        return rays_o, rays_v
//  [python]
def gen_rays_at(self, img_idx, resolution_level=1):
        """
        Generate rays at world space from one camera.
        """
        l = resolution_level
        tx = torch.linspace(0, self.W - 1, int(self.W // l))
        ty = torch.linspace(0, self.H - 1, int(self.H // l))
        pixels_x, pixels_y = torch.meshgrid(tx, ty)
        pixels_x = pixels_x.t()
        pixels_y = pixels_y.t()
        p = torch.stack([(pixels_x - self.K[0][2]) / self.K[0][0],
                         -(pixels_y - self.K[1][2]) / self.K[1][1],
                         -torch.ones_like(pixels_x)], -1).float()
        rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # W, H, 3
        # rays_v = torch.matmul(self.poses[img_idx, None, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # W, H, 3
        rays_v = torch.sum(rays_v[..., None, :] * self.poses[img_idx, :3, :3], -1)
        rays_o = self.poses[img_idx, None, None, :3, 3].expand(rays_v.shape)  # W, H, 3
        return rays_o, rays_v
//  [python]
def writepdb(filename, atoms, seq, Ls, idx_pdb=None, bfacts=None):
    f = open(filename,"w")
    ctr = 1
    scpu = seq.cpu().squeeze(0)
    atomscpu = atoms.cpu().squeeze(0)

    if bfacts is None:
        bfacts = torch.zeros(atomscpu.shape[0])
    if idx_pdb is None:
        idx_pdb = 1 + torch.arange(atomscpu.shape[0])

    Bfacts = torch.clamp( bfacts.cpu(), 0, 100)
    chn_idx, res_idx = 0, 0
    for i,s in enumerate(scpu):
        natoms = atomscpu.shape[-2]
        if (natoms!=NHEAVY and natoms!=NTOTAL):
            print ('bad size!', natoms, NHEAVY, NTOTAL, atoms.shape)
            assert(False)

        atms = aa2long[s]

        # his protonation state hack
        if (s==8 and torch.norm( atomscpu[i,9,:]-atomscpu[i,5,:] ) < 1.7):
            atms = (
                " N  "," CA "," C  "," O  "," CB "," CG "," NE2"," CD2"," CE1"," ND1",
                  None,  None,  None,  None," H  "," HA ","1HB ","2HB "," HD2"," HE1",
                " HD1",  None,  None,  None,  None,  None,  None) # his_d

        for j,atm_j in enumerate(atms):
            if (j<natoms and atm_j is not None and not torch.isnan(atomscpu[i,j,:]).any()):
                f.write ("%-6s%5s %4s %3s %s%4d    %8.3f%8.3f%8.3f%6.2f%6.2f\n"%(
                    "ATOM", ctr, atm_j, num2aa[s], 
                    PDB_CHAIN_IDS[chn_idx], res_idx+1, atomscpu[i,j,0], atomscpu[i,j,1], atomscpu[i,j,2],
                    1.0, Bfacts[i] ) )
                ctr += 1

        res_idx += 1
        if (chn_idx < len(Ls) and res_idx == Ls[chn_idx]):
            chn_idx += 1
            res_idx = 0
//  [python]
def writepdb(filename, atoms, seq, Ls, idx_pdb=None, bfacts=None):
    f = open(filename,"w")
    ctr = 1
    scpu = seq.cpu().squeeze(0)
    atomscpu = atoms.cpu().squeeze(0)

    if bfacts is None:
        bfacts = torch.zeros(atomscpu.shape[0])
    if idx_pdb is None:
        idx_pdb = 1 + torch.arange(atomscpu.shape[0])

    Bfacts = torch.clamp( bfacts.cpu(), 0, 100)
    chn_idx, res_idx = 0, 0
    for i,s in enumerate(scpu):
        natoms = atomscpu.shape[-2]
        if (natoms!=NHEAVY and natoms!=NTOTAL):
            print ('bad size!', natoms, NHEAVY, NTOTAL, atoms.shape)
            assert(False)

        atms = aa2long[s]

        # his protonation state hack
        if (s==8 and torch.linalg.norm( atomscpu[i,9,:]-atomscpu[i,5,:] ) < 1.7):
            atms = (
                " N  "," CA "," C  "," O  "," CB "," CG "," NE2"," CD2"," CE1"," ND1",
                  None,  None,  None,  None," H  "," HA ","1HB ","2HB "," HD2"," HE1",
                " HD1",  None,  None,  None,  None,  None,  None) # his_d

        for j,atm_j in enumerate(atms):
            if (j<natoms and atm_j is not None and not torch.isnan(atomscpu[i,j,:]).any()):
                f.write ("%-6s%5s %4s %3s %s%4d    %8.3f%8.3f%8.3f%6.2f%6.2f\n"%(
                    "ATOM", ctr, atm_j, num2aa[s], 
                    PDB_CHAIN_IDS[chn_idx], res_idx+1, atomscpu[i,j,0], atomscpu[i,j,1], atomscpu[i,j,2],
                    1.0, Bfacts[i] ) )
                ctr += 1

        res_idx += 1
        if (chn_idx < len(Ls) and res_idx == Ls[chn_idx]):
            chn_idx += 1
            res_idx = 0
//  [python]
def gen_rays_silhouettes(self, pose, max_ray_num, mask):
        if mask.sum() == 0:
            return self.gen_rays_pose(pose, resolution_level=4)
        struct = ndimage.generate_binary_structure(2, 2)
        dilated_mask = ndimage.binary_dilation(mask, structure=struct, iterations=10).astype(np.int32)
        current_ratio = dilated_mask.sum() / float(mask.shape[0] * mask.shape[1])
        W = H = min(self.H, int(np.sqrt(max_ray_num / current_ratio)))
        tx = torch.linspace(0, self.W - 1, W)
        ty = torch.linspace(0, self.H - 1, H)
        pixels_x, pixels_y = torch.meshgrid(tx, ty)
        pixels_x = pixels_x.t()
        pixels_y = pixels_y.t()
        p = torch.stack([(pixels_x - self.K[0][2]) / self.K[0][0],
                         -(pixels_y - self.K[1][2]) / self.K[1][1],
                         -torch.ones_like(pixels_x)], -1).float()
        rays_v = p / torch.norm(p, p=2, dim=-1, keepdim=True)  # W, H, 3
        rays_v = torch.sum(rays_v[..., None, :] * pose[:3, :3], -1)
        rays_o = pose[None, None, :3, 3].expand(rays_v.shape)  # W, H, 3
        resized_dilated_mask = torch.nn.functional.interpolate(
            torch.from_numpy(dilated_mask).reshape(256, 256, 1).permute(2, 0, 1).unsqueeze(0).float(), size=(H, W)).squeeze()
        masked_rays_v = rays_v[resized_dilated_mask > 0]
        masked_rays_o = rays_o[resized_dilated_mask > 0]

        return masked_rays_o, masked_rays_v, W, resized_dilated_mask > 0
//  [python]
def gen_rays_silhouettes(self, pose, max_ray_num, mask):
        if mask.sum() == 0:
            return self.gen_rays_pose(pose, resolution_level=4)
        struct = ndimage.generate_binary_structure(2, 2)
        dilated_mask = ndimage.binary_dilation(mask, structure=struct, iterations=10).astype(np.int32)
        current_ratio = dilated_mask.sum() / float(mask.shape[0] * mask.shape[1])
        W = H = min(self.H, int(np.sqrt(max_ray_num / current_ratio)))
        tx = torch.linspace(0, self.W - 1, W)
        ty = torch.linspace(0, self.H - 1, H)
        pixels_x, pixels_y = torch.meshgrid(tx, ty)
        pixels_x = pixels_x.t()
        pixels_y = pixels_y.t()
        p = torch.stack([(pixels_x - self.K[0][2]) / self.K[0][0],
                         -(pixels_y - self.K[1][2]) / self.K[1][1],
                         -torch.ones_like(pixels_x)], -1).float()
        rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # W, H, 3
        rays_v = torch.sum(rays_v[..., None, :] * pose[:3, :3], -1)
        rays_o = pose[None, None, :3, 3].expand(rays_v.shape)  # W, H, 3
        resized_dilated_mask = torch.nn.functional.interpolate(
            torch.from_numpy(dilated_mask).reshape(256, 256, 1).permute(2, 0, 1).unsqueeze(0).float(), size=(H, W)).squeeze()
        masked_rays_v = rays_v[resized_dilated_mask > 0]
        masked_rays_o = rays_o[resized_dilated_mask > 0]

        return masked_rays_o, masked_rays_v, W, resized_dilated_mask > 0
//  [python]
def get_ray_directions(W: int, H: int, fx: float, fy: float, cx: float, cy: float, center_pixels: bool,
                       device: torch.device) -> torch.Tensor:
    i, j = torch.meshgrid(torch.arange(W, dtype=torch.float32, device=device),
                          torch.arange(H, dtype=torch.float32, device=device), indexing='xy')
    if center_pixels:
        i = i.clone() + 0.5
        j = j.clone() + 0.5

    directions = \
        torch.stack([(i - cx) / fx, -(j - cy) / fy, -torch.ones_like(i)], -1)  # (H, W, 3)
    directions /= torch.norm(directions, dim=-1, keepdim=True)

    return directions
//  [python]
def get_ray_directions(W: int, H: int, fx: float, fy: float, cx: float, cy: float, center_pixels: bool,
                       device: torch.device) -> torch.Tensor:
    i, j = torch.meshgrid(torch.arange(W, dtype=torch.float32, device=device),
                          torch.arange(H, dtype=torch.float32, device=device), indexing='xy')
    if center_pixels:
        i = i.clone() + 0.5
        j = j.clone() + 0.5

    directions = \
        torch.stack([(i - cx) / fx, -(j - cy) / fy, -torch.ones_like(i)], -1)  # (H, W, 3)
    directions /= torch.linalg.norm(directions, dim=-1, keepdim=True)

    return directions
//  [python]
def gen_random_rays_at(self, img_idx, batch_size):
        """
        Generate random rays at world space from one camera.
        """
        pixels_x = torch.randint(low=0, high=self.W, size=[batch_size])
        pixels_y = torch.randint(low=0, high=self.H, size=[batch_size])
        color = self.images[img_idx][(pixels_y, pixels_x)]    # batch_size, 3
        mask = self.masks[img_idx][(pixels_y, pixels_x)]      # batch_size, 3
        p = torch.stack([(pixels_x - self.K[0][2]) / self.K[0][0],
                         -(pixels_y - self.K[1][2]) / self.K[1][1],
                         -torch.ones_like(pixels_x)], -1).float()
        rays_v = p / torch.norm(p, p=2, dim=-1, keepdim=True)    # batch_size, 3
        # rays_v = torch.matmul(self.poses[img_idx, None, :3, :3], rays_v[:, :, None]).squeeze()  # batch_size, 3
        rays_v = torch.sum(rays_v[..., None, :] * self.poses[img_idx, :3, :3], -1)
        rays_o = self.poses[img_idx, None, :3, 3].expand(rays_v.shape) # batch_size, 3
        return torch.cat([rays_o.cpu(), rays_v.cpu(), color, mask[:, :1]], dim=-1).cuda()