code_snapshots/radiomap_compare_code/NeRF2/renderer.py

# -*- coding: utf-8 -*-
"""code for ray marching and signal rendering
"""
import torch
import numpy as np
import torch.nn.functional as F
import scipy.constants as sc
from einops import rearrange, repeat


class Renderer():

    def __init__(self, networks_fn, **kwargs) -> None:
        """
        Parameters
        -----------------
        near : float. The near bound of the rays
        far : float. The far bound of the rays
        n_samples: int. num of samples per ray
        """

        ## Rendering parameters
        self.network_fn = networks_fn
        self.n_samples = kwargs['n_samples']
        self.near = kwargs['near']
        self.far = kwargs['far']


    def sample_points(self, rays_o, rays_d):
        """sample points along rays

        Parameters
        ----------
        rays_o : tensor. [n_rays, 3]. The origin of rays
        rays_d : tensor. [n_rays, 3]. The direction of rays

        Returns
        -------
        pts : tensor. [n_rays, n_samples, 3]. The sampled points along rays
        t_vals : tensor. [n_rays, n_samples]. The distance from origin to each sampled point
        """
        shape = list(rays_o.shape)
        shape[-1] = 1
        near, far = torch.full(shape, self.near), torch.full(shape, self.far)
        t_vals = torch.linspace(0., 1., steps=self.n_samples) * (far - near) + near  # scale t with near and far
        t_vals = t_vals.to(rays_o.device)
        pts = rays_o[...,None,:] + rays_d[...,None,:] * t_vals[...,:,None] # p = o + td, [n_rays, n_samples, 3]

        return pts, t_vals




class Renderer_spectrum(Renderer):
    """Renderer for spectrum (integral from single direction)
    """

    def __init__(self, networks_fn, **kwargs) -> None:
        """
        Parameters
        -----------------
        near : float. The near bound of the rays
        far : float. The far bound of the rays
        n_samples: int. num of samples per ray
        """
        super().__init__(networks_fn, **kwargs)



    def render_ss(self, tx, rays_o, rays_d):
        """render the signal strength of each ray

        Parameters
        ----------
        tx: tensor. [batchsize, 3]. The position of the transmitter
        rays_o : tensor. [batchsize, 3]. The origin of rays
        rays_d : tensor. [batchsize, 3]. The direction of rays
        """

        # sample points along rays
        pts, t_vals = self.sample_points(rays_o, rays_d)

        # Expand views and tx to match the shape of pts
        view = rays_d[:, None].expand(pts.shape)
        tx = tx[:, None].expand(pts.shape)

        # Run network and compute outputs
        raw = self.network_fn(pts, view, tx)    # [batchsize, n_samples, 4]
        receive_ss = self.raw2outputs(raw, t_vals, rays_d)  # [batchsize]

        return receive_ss



    def raw2outputs(self, raw, r_vals, rays_d):
        """Transforms model's predictions to semantically meaningful values. (core part)

        Parameters
        ----------
        raw : [batchsize, n_samples, 4]. Prediction from model.
        r_vals : [batchsize, n_samples]. Integration distance.
        rays_d : [batchsize, 3]. Direction of each ray

        Return:
        ----------
        receive_signal : [batchsize]. abs(singal of each ray)
        """

        raw2alpha = lambda raw, dists: 1.-torch.exp(-raw*dists)
        # raw2phase = lambda raw, dists: torch.exp(1j*raw*dists)
        raw2phase = lambda raw, dists: raw*dists

        dists = r_vals[...,1:] - r_vals[...,:-1]
        dists = torch.cat([dists, torch.Tensor([1e10]).cuda().expand(dists[...,:1].shape)], -1)  # [N_rays, N_samples]
        dists = dists * torch.norm(rays_d[...,None,:], dim=-1)

        att_a, att_p, s_a, s_p = raw[...,0], raw[...,1], raw[...,2], raw[...,3]    # [N_rays, N_samples]
        att_p, s_p = torch.sigmoid(att_p)*np.pi*2, torch.sigmoid(s_p)*np.pi*2
        att_a, s_a = abs(F.leaky_relu(att_a)), abs(F.leaky_relu(s_a))
        # att_a, s_a = torch.sigmoid(att_a), torch.sigmoid(s_a)

        alpha = raw2alpha(att_a, dists)  # [N_rays, N_samples]
        phase = raw2phase(att_p, dists)

        att_i = torch.cumprod(torch.cat([torch.ones((alpha.shape[0], 1)).cuda(), 1.-alpha + 1e-10], -1), -1)[:, :-1]
        path = torch.cat([r_vals[...,1:], torch.Tensor([1e10]).cuda().expand(r_vals[...,:1].shape)], -1)
        path_loss = 0.025 / path
        # phase_i = torch.cumprod(torch.cat([torch.ones((alpha.shape[0], 1)), phase], -1), -1)[:, :-1]
        phase_i = torch.cumsum(torch.cat([torch.ones((alpha.shape[0], 1)).cuda(), phase], -1), -1)[:, :-1]
        phase_i = torch.exp(1j*phase_i)    # [N_rays, N_samples]


        receive_signal = torch.sum(s_a*torch.exp(1j*s_p)*att_i*phase_i*path_loss, -1)  # [N_rays]
        receive_signal = abs(receive_signal)

        return receive_signal




class Renderer_RSSI(Renderer):
    """Renderer for RSSI (integral from all directions)
    """

    def __init__(self, networks_fn, **kwargs) -> None:
        """
        Parameters
        -----------------
        near : float. The near bound of the rays
        far : float. The far bound of the rays
        n_samples: int. num of samples per ray
        """
        super().__init__(networks_fn, **kwargs)


    def render_rssi(self, tx, rays_o, rays_d):
        """render the RSSI for each gateway. To avoid OOM, we split the rays into chunks

        Parameters
        ----------
        tx: tensor. [batchsize, 3]. The position of the transmitter
        rays_o : tensor. [batchsize, 3]. The origin of rays
        rays_d : tensor. [batchsize, 9x36x3]. The direction of rays
        """

        batchsize, _ = tx.shape
        rays_d = torch.reshape(rays_d, (batchsize, -1, 3))    # [batchsize, 9x36, 3]
        chunks = 36
        chunks_num = 36 // chunks
        rays_o_chunk = rays_o.expand(chunks, -1, -1).permute(1,0,2) #[bs, cks, 3]
        tags_chunk = tx.expand(chunks, -1, -1).permute(1,0,2)        #[bs, cks, 3]
        recv_signal = torch.zeros(batchsize).cuda()
        for i in range(chunks_num):
            rays_d_chunk = rays_d[:,i*chunks:(i+1)*chunks, :]  # [bs, cks, 3]
            pts, t_vals = self.sample_points(rays_o_chunk, rays_d_chunk) # [bs, cks, pts, 3]
            views_chunk = rays_d_chunk[..., None, :].expand(pts.shape)  # [bs, cks, pts, 3]
            tx_chunk = tags_chunk[..., None, :].expand(pts.shape)  # [bs, cks, pts, 3]

            # Run network and compute outputs
            raw = self.network_fn(pts, views_chunk, tx_chunk)    # [batchsize, chunks, n_samples, 4]
            recv_signal_chunks = self.raw2outputs_signal(raw, t_vals, rays_d_chunk)  # [bs]
            recv_signal += recv_signal_chunks

        return recv_signal    # [batchsize,]


    def raw2outputs_signal(self, raw, r_vals, rays_d):
        """Transforms model's predictions to semantically meaningful values.

        Parameters
        ----------
        raw : [batchsize, chunks,n_samples,  4]. Prediction from model.
        r_vals : [batchsize, chunks, n_samples]. Integration distance.
        rays_d : [batchsize,chunks, 3]. Direction of each ray

        Return:
        ----------
        receive_signal : [batchsize]. abs(singal of each ray)
        """
        wavelength = sc.c / 2.4e9
        # raw2alpha = lambda raw, dists: 1.-torch.exp(-raw*dists)
        # raw2phase = lambda raw, dists: torch.exp(1j*raw*dists)
        raw2phase = lambda raw, dists: raw + 2*np.pi*dists/wavelength
        raw2amp = lambda raw, dists: -raw*dists

        dists = r_vals[...,1:] - r_vals[...,:-1]
        dists = torch.cat([dists, torch.Tensor([1e10]).cuda().expand(dists[...,:1].shape)], -1)  # [batchsize, chunks, n_samples]
        dists = dists * torch.norm(rays_d[...,None,:], dim=-1)  # [batchsize,chunks, n_samples, 3].

        att_a, att_p, s_a, s_p = raw[...,0], raw[...,1], raw[...,2], raw[...,3]    # [batchsize,chunks, N_samples]
        att_p, s_p = torch.sigmoid(att_p)*np.pi*2-np.pi, torch.sigmoid(s_p)*np.pi*2-np.pi
        att_a, s_a = abs(F.leaky_relu(att_a)), abs(F.leaky_relu(s_a))

        amp = raw2amp(att_a, dists)  # [batchsize,chunks, N_samples]
        phase = raw2phase(att_p, dists)

        # att_i = torch.cumprod(torch.cat([torch.ones((al_shape[:-1], 1)), 1.-alpha + 1e-10], -1), -1)[:, :-1]
        # phase_i = torch.cumprod(torch.cat([torch.ones((alpha.shape[0], 1)), phase], -1), -1)[:, :-1]
        amp_i = torch.exp(torch.cumsum(amp, -1))            # [batchsize,chunks, N_samples]
        phase_i = torch.exp(1j*torch.cumsum(phase, -1))                # [batchsize,chunks, N_samples]

        recv_signal = torch.sum(s_a*torch.exp(1j*s_p)*amp_i*phase_i, -1)  # integral along line [batchsize,chunks]
        recv_signal = torch.sum(recv_signal, -1)   # integral along direction [batchsize,]

        return abs(recv_signal)




class Renderer_CSI(Renderer):
    """Renderer for CSI (integral from all directions)
    """

    def __init__(self, networks_fn, **kwargs) -> None:
        """
        Parameters
        -----------------
        near : float. The near bound of the rays
        far : float. The far bound of the rays
        n_samples: int. num of samples per ray
        """
        super().__init__(networks_fn, **kwargs)



    def render_csi(self, uplink, rays_o, rays_d):
        """render the RSSI for each gateway.

        Parameters
        ----------
        uplink: tensor. [batchsize, 52]. The uplink CSI (26 real + 26 imag)
        rays_o : tensor. [batchsize, 3]. The origin of rays
        rays_d : tensor. [batchsize, 9x36x3]. The direction of rays
        """

        rays_d = rearrange(rays_d, 'b (v d) -> b v d', d=3)    # [bs, 9x36, 3]
        batchsize, viewsize, _ = rays_d.shape
        rays_o = repeat(rays_o, 'b d -> b v d', v=viewsize) #[bs, 9x36, 3]
        uplink = repeat(uplink, 'b d -> b v d', v=viewsize)        #[bs, 9x36, 52]

        pts, t_vals = self.sample_points(rays_o, rays_d) # [bs, 9x36, pts, 3]
        views = repeat(rays_d, 'b v d -> b v p d', p=self.n_samples)  # [bs, 9x36, pts, 3]
        uplink = repeat(uplink, 'b v d -> b v p d', p=self.n_samples)  # [bs, cks, pts, 3]

        # Run network and compute outputs
        raw = self.network_fn(pts, views, uplink)    # [batchsize, 9x36, pts, 4]
        recv_signal = self.raw2outputs_signal(raw, t_vals, rays_d)  # [bs, 26]

        return recv_signal


    def raw2outputs_signal(self, raw, r_vals, rays_d):
        """Transforms model's predictions to semantically meaningful values.

        Parameters
        ----------
        raw : [batchsize, chunks,n_samples,  4]. Prediction from model.
        r_vals : [batchsize, chunks, n_samples]. Integration distance.
        rays_d : [batchsize,chunks, 3]. Direction of each ray

        Return:
        ----------
        receive_signal : [batchsize, 26]. OFDM singal of each ray
        """
        wavelength = sc.c / 2.4e9
        # raw2alpha = lambda raw, dists: 1.-torch.exp(-raw*dists)
        # raw2phase = lambda raw, dists: torch.exp(1j*raw*dists)
        raw2phase = lambda raw, dists: raw + 2*np.pi*dists/wavelength
        raw2amp = lambda raw, dists: -raw*dists

        dists = r_vals[...,1:] - r_vals[...,:-1]
        dists = torch.cat([dists, torch.Tensor([1e10]).cuda().expand(dists[...,:1].shape)], -1)  # [batchsize, chunks, n_samples]
        dists = dists * torch.norm(rays_d[...,None,:], dim=-1)  # [batchsize,chunks, n_samples, 3].

        att_a, att_p, s_a, s_p = raw[...,:26], raw[...,26:52], raw[...,52:78], raw[...,78:104]    # [batchsize,chunks, N_samples]
        att_p, s_p = torch.sigmoid(att_p)*np.pi*2-np.pi, torch.sigmoid(s_p)*np.pi*2-np.pi
        att_a, s_a = abs(F.leaky_relu(att_a)), abs(F.leaky_relu(s_a))

        dists = dists.unsqueeze(-1)

        amp = raw2amp(att_a, dists)  # [batchsize,chunks, N_samples, 26]
        phase = raw2phase(att_p, dists)

        # att_i = torch.cumprod(torch.cat([torch.ones((al_shape[:-1], 1)), 1.-alpha + 1e-10], -1), -1)[:, :-1]
        # phase_i = torch.cumprod(torch.cat([torch.ones((alpha.shape[0], 1)), phase], -1), -1)[:, :-1]
        amp_i = torch.exp(torch.cumsum(amp, -2))            # [batchsize,chunks, N_samples, 26]
        phase_i = torch.exp(1j*torch.cumsum(phase, -2))                # [batchsize,chunks, N_samples 26]

        recv_signal = torch.sum(s_a*torch.exp(1j*s_p)*amp_i*phase_i, -2)  # integral along line [batchsize,chunks,26]
        recv_signal = torch.sum(recv_signal, 1)   # integral along direction [batchsize, 26]

        return recv_signal




renderer_dict = {"spectrum": Renderer_spectrum, "rssi": Renderer_RSSI, "csi": Renderer_CSI}