Method forward

locationencoder/pe/grid_and_sphere.py:94–247

(self, coords)

        self.freq_mat = np.repeat(freq_mat, 2, axis=1)

    def forward(self, coords):
        device = coords.device
        dtype = coords.dtype
        N = coords.size(0)

        # add 1 context point dimension (unused here)
        coords = coords[:, None, :]

        # coords_mat: shape (batch_size, num_context_pt, 2)
        coords_mat = np.asarray(coords.cpu())
        batch_size = coords_mat.shape[0]
        num_context_pt = coords_mat.shape[1]
        # coords_mat: shape (batch_size, num_context_pt, 2, 1)
        coords_mat = np.expand_dims(coords_mat, axis=3)
        # coords_mat: shape (batch_size, num_context_pt, 2, 1, 1)
        coords_mat = np.expand_dims(coords_mat, axis=4)
        # coords_mat: shape (batch_size, num_context_pt, 2, frequency_num, 1)
        coords_mat = np.repeat(coords_mat, self.frequency_num, axis=3)
        # coords_mat: shape (batch_size, num_context_pt, 2, frequency_num, 2)
        coords_mat = np.repeat(coords_mat, 2, axis=4)
        # spr_embeds: shape (batch_size, num_context_pt, 2, frequency_num, 2)
        spr_embeds = coords_mat * self.freq_mat

        if self.name == "grid":
            # eq 3 in https://arxiv.org/pdf/2201.10489.pdf
            # code from https://github.com/gengchenmai/space2vec/blob/a29793336e6a1ebdb497289c286a0b4d5a83079f/spacegraph/spacegraph_codebase/SpatialRelationEncoder.py#L135

            spr_embeds[:, :, :, :, 0::2] = np.sin(spr_embeds[:, :, :, :, 0::2])  # dim 2i
            spr_embeds[:, :, :, :, 1::2] = np.cos(spr_embeds[:, :, :, :, 1::2])  # dim 2i+1

        elif self.name == "spherec":
            # eq 4 in https://arxiv.org/pdf/2201.10489.pdf
            # lambda: longitude, theta=latitude

            #sin_lon, sin_lat = np.sin(spr_embeds[:, 0, :, :, 0]).transpose(1, 0, 2)
            #cos_lon, cos_lat = np.cos(spr_embeds[:, 0, :, :, 1]).transpose(1, 0, 2)

            # eq 4
            # sin theta, cos_theta * cos_lambda, cos_theta * sin_lambda
            # sin lat, cos_lat cos_lon, cos_lat sin_lon
            #spr_embeds = np.stack([sin_lat, cos_lat*cos_lon, cos_lat*sin_lon], axis=-1)

            spr_embeds = spr_embeds# * math.pi / 180

            # lon, lat: shape (batch_size, num_context_pt, 1, frequency_num, 1)
            lon = np.expand_dims(spr_embeds[:, :, 0, :, :], axis=2)
            lat = np.expand_dims(spr_embeds[:, :, 1, :, :], axis=2)

            # make sinuniod function
            # lon_sin, lon_cos: shape (batch_size, num_context_pt, 1, frequency_num, 1)
            lon_sin = np.sin(lon)
            lon_cos = np.cos(lon)

            # lat_sin, lat_cos: shape (batch_size, num_context_pt, 1, frequency_num, 1)
            lat_sin = np.sin(lat)
            lat_cos = np.cos(lat)

            # spr_embeds_: shape (batch_size, num_context_pt, 1, frequency_num, 3)