LG_generation_GPU.py - dev - stronius/OAM_ROM_Model

LG_generation_GPU.py
522 строки18 KB
stronius
Readme update, Example minor vis update, code harm_cutter and oam_spectrum update
10 фев 2025, 15:12
10 фев 2025, 15:12a2fbd82
from datetime import datetime
import numpy as np
import cupy as cp
from numba import jit, float64, complex128, int64, types
from scipy.special import genlaguerre, eval_genlaguerre

from dataclasses import dataclass

from typing import List, Tuple, Literal
from math import factorial

from scipy.stats import trim_mean

@dataclass(frozen=True)
class BeamConfig:
    """
    Class for configuration storage. Information about computational box, beams and delays stores here.

    Attrs:
        modes (List[Tuple[float, int, int, int, float, float]]): 
            Each mode description in format (a0, m, p, Mp, w0, tau).
        delays (List[float]): 
            Delays between modes relative to z = 0.
        coord_type (Literal["cartesian", "polar"]): 
            Type of coordinate system ('cartesian' или 'polar'). Dy default 'polar'.
        limits (Tuple[float, float, float, float, float, float]): 
            Limits for computational box 
            For 'cartesian': (xmin, xmax, ymin, ymax, zmin, zmax). 
            For 'polar': (rhomin, rhomax, phimin, phimax, zmin, zmax).
        discretization (Tuple[float, float, float]): 
            Coordinates discretization: number of points. 
            For 'cartesian': (nx, ny, nz). 
            For 'polar': (nrho, nphi, nz).
    """
    
    modes: List[Tuple[float, int, int, int, float, float]]
    delays: List[float]
    coord_type: Literal["cartesian", "polar"]
    limits: Tuple[float, float, float, float, float, float]
    discretization: Tuple[float, float, float]

    def __post_init__(self):
        """Check correctness after __init__"""
        self._validate_config()

    def _validate_config(self):
        """Check correctness of the parameters"""
        if self.coord_type == "cartesian":
            xmin, xmax, ymin, ymax, zmin, zmax = self.limits
            assert xmin < xmax, "xmin should be less than xmax"
            assert ymin < ymax, "ymin should be less than ymax"
            assert zmin < zmax, "zmin should be less than zmax"
        elif self.coord_type == "polar":
            rhomin, rhomax, phimin, phimax, zmin, zmax = self.limits
            assert 0 <= rhomin < rhomax, "rhomin should be less than rhomax and greater than 0"
            assert phimin < phimax, "phimin should be less than phimax"
            assert zmin < zmax, "zmin should be less than zmax"
        else:
            raise ValueError("Coordinate type should be 'cartesian' or 'polar'")

        assert len(self.delays) == len(self.modes), "Number of delays should equal to number of modes"
        dx, dy, dz = self.discretization
        assert dx > 0 and dy > 0 and dz > 0, "Discretization number of points should be greater than 0"


def create_and_validate_config(
    modes: List[Tuple[float, int, int, int, float, float]] = [(1, 0, 0, 1, 3*2*np.pi, 3*2*np.pi)],
    delays: List[float] = [-2*np.pi],
    coord_type: Literal["cartesian", "polar"] = "polar",
    limits: Tuple[float, float, float, float, float, float] = (0, 7, 0.0, 2*np.pi, -15.0, 15.0),
    discretization: Tuple[float, float, float] = (20, 150, 5000),
    verbose: bool = False
    ) -> BeamConfig:
    """
    Function for Creation and validation of the configuration

    Args:
        modes (List[Tuple[float, int, int, int, float, float]]): 
            Mode description: A0, l_z, p, M_p, W0, tau. By default Gaussian Beam [(1, 0, 0, 1, 3*2*np.pi, 3*2*np.pi)].
        delays (List[float]): 
            Delays between modes. By default [-2*np.pi].
        coord_type (Literal["cartesian", "polar"]): 
            Coordinate system type. By default 'polar'.
        limits (Tuple[float, float, float, float, float, float]): 
            Limits for computational box. 
            By default for 'polar': (0, 7, 0.0, 2*np.pi, -15.0, 15.0).
            Please, set the limit implicitly w.r.t. wavelength. All spatial values will be multiplied by 2*pi, except phi angle. 
        discretization (Tuple[float, float, float]): 
            Discretization number of points. By default (20, 150, 5000).
        verbose bool:
            Print configuration info.  

    Returns:
        GaussianBeamConfig: configuration object.

    Exceptions:
        AssertionError: If parameters are incorrect.
        ValueError: If parameters are invalid.
    """
    try:
        config = BeamConfig(
            modes=modes,
            delays=delays,
            coord_type=coord_type,
            limits=limits,
            discretization=discretization,
        )
        if verbose:
            print("Configuration creation Success:", config, sep='\n')
        else:
            print("Configuration creation Success")
        return config
    except (AssertionError, ValueError) as e:
        print(f"Configuration Error: {e}")
        raise
        
class LG_mode_generation_GPU():
    
     
    """
    Class for Model calculation. 
    
    Attrs:
        config - configuration object as BeamConfig. If not specified - default config will be generated with Gaussian beam.
    """
    
    def __init__(self, configuration=None):
        
        self.fitted = False
        self.spectrated = False
        
        if configuration is None:
            self.config = create_and_validate_config()
        else:
            self.config = configuration
            
        if self.config.coord_type == "cartesian":
            
            xmin, xmax, ymin, ymax, zmin, zmax = self.config.limits
            nx, ny, nz = self.config.discretization
            
            xx = np.linspace(xmin*2*np.pi, xmax*2*np.pi, nx)
            yy = np.linspace(ymin*2*np.pi, ymax*2*np.pi, ny)
            zz = np.linspace(zmin*2*np.pi, zmax*2*np.pi, nz)
            
            self.zmin, self.zmax = zmin, zmax
            self.d_sampling_rate = zz[1] - zz[0]
            self.nz = nz
            self.x,self.y,self.z = np.meshgrid(xx,yy,zz, indexing='ij')
            
        elif self.config.coord_type == "polar":
            
            rhomin, rhomax, phimin, phimax, zmin, zmax = self.config.limits
            nrho, nphi, nz = self.config.discretization
            
            zz = np.linspace(zmin*2*np.pi, zmax*2*np.pi, nz)
            rhos = np.linspace(rhomin*2*np.pi, rhomax*2*np.pi, nrho)
            phis = np.linspace(phimin, phimax, nphi)
            self.z_grid = zz.copy()
            
            self.d_sampling_rate = zz[1] - zz[0]
            self.nz, self.nrho, self.nphi = nz, nrho, nphi
            self.zmin, self.zmax = zmin, zmax
            
            self.rho,self.phi,self.z = np.meshgrid(rhos,phis,zz, indexing='ij')
            
    @staticmethod
    def LG_mode_GPU(r, phi, z, mode_params_a0, mode_params_m, 
                        mode_params_p, 
                        mode_params_Mp, 
                        mode_params_w0, 
                        tau):

        """
        Calculates LG mode field with GPU acceleration.

        Parameters:
        ----------
        r : float64[:,:,:]
            Radials in polar coordinates.
        phi : float64[:,:,:]
            Angle in polar coordinates.
        z : float64[:,:,:]
            Propagation axis.
        mode_params_a0 : float64
            Amplitude.
        mode_params_m : int64
            Mode number OAM.
        mode_params_p : int64
            Mode number radial momentum.
        mode_params_Mp : float64
            Frequency parameter.
        mode_params_w0 : float64
            Beam waist.
        tau : float64
            Beam width.

        Returns:
        -----------
        tuple
            field, argument for Generalizes LG polynomial, OAM, radial momentum
        """
        
        a0 = mode_params_a0
        m = mode_params_m  
        p = mode_params_p  
        Mp = mode_params_Mp
        w0 = mode_params_w0
    
        Z_R = Mp * w0**2 / 2
        
        r_sq = cp.square(r)
        z_sq = cp.square(z)
        
        w = w0 * cp.sqrt(1 + z_sq / Z_R**2)
        w_sq = cp.square(w)
        R_z = z * (1 + Z_R**2 / z_sq)
        chi = (2 * p + np.abs(m) + 1) * cp.arctan(z / Z_R)
        
        toret = (w0 / w) * (r * cp.sqrt(2) / w)**(cp.abs(m)) * cp.exp(-r_sq / w_sq)
        toret = toret * cp.exp(-1j * Mp * z + 1j * m * phi + 1j * r_sq / 2 / R_z - 1j*chi - z_sq / 2 / tau**2)

        Lpm_values_cpu = eval_genlaguerre(p, np.abs(m), (2 * r_sq / w_sq).get())

        Lpm_values = cp.asarray(Lpm_values_cpu)
        coeff = cp.sqrt(2 * factorial(p) / (np.pi * factorial(abs(m) + p)))
        
        res = a0 * toret * coeff * Lpm_values
    
        return res         
        
    def cartesian_to_polar(self):
        """
        Function transforms cartesian coordinates from config and make is polar. 
        Returns: rho, phi, z arrays
        """
        if self.config.coord_type =="polar":
            return self.rho, self.phi, self.z
        else:
            r = np.sqrt(self.x**2+self.y**2)

            xiy = self.x + 1j*self.y
            phi = np.angle(xiy)
            return r, phi, self.z
        
    def polar_to_cartesian(self):
        ...
    
    def calculate_field(self):
        """ 
        Function Calculates the incident field and reflected field based on config.
        Beams are allocated according to delays. 
        Only cosine beams are supported yet.
        After this, the following results are done:
            Incident field
            Reflected field
            zmirror
            full_envelope
        To proceed, call calculate_spectrum function. 
        """
        start = datetime.now()
        start_event = cp.cuda.Event()
        copy_in_event = cp.cuda.Event()
        inc_event = cp.cuda.Event()
        refl_event = cp.cuda.Event()
        interp_event = cp.cuda.Event()
        copy_event = cp.cuda.Event()
        
        print("Calculating LG modes...", end="")
        
        start_event.record()
        
        
        r, phi, z = self.cartesian_to_polar()
        r_d, phi_d, z_d = cp.asarray(r), cp.asarray(phi), cp.asarray(z)
        
        copy_in_event.record()
        copy_in_event.synchronize()
        
        modes, delays = self.config.modes, self.config.delays
        inc_field = 0
        cos_sum = 0
        sin_sum = 0
        LG_envs = [0] * len(modes)
        refl_field = 0
        
        for i, (mode, delay) in enumerate(zip(modes, delays)):
                   
            LG = self.LG_mode_GPU(r_d, phi_d, z_d-delay, *mode)

            LG_abs = cp.abs(LG)
            LG_angle = cp.angle(LG)
            
            inc_field = cp.add(inc_field, cp.multiply(LG_abs, cp.cos(LG_angle)))
            cos_sum = cp.add(cos_sum, inc_field)
            sin_sum = cp.add(sin_sum, cp.multiply(LG_abs, cp.sin(LG_angle)))
            
            LG_envs[i] = LG_abs
                    
        full_envelope = cp.sqrt(cp.add(cp.square(cos_sum),cp.square(sin_sum**2)))
        
        inc_event.record()
        inc_event.synchronize()
        
        
        zmirror_d = cos_sum**2 / (1 + full_envelope**2)
        
        for i, (mode, delay) in enumerate(zip(modes, delays)):
            
            LG = self.LG_mode_GPU(r_d, phi_d, z_d-delay - zmirror_d, *mode)

            LG_angle = cp.angle(LG)
           
            refl_field = cp.add(refl_field, cp.multiply(LG_envs[i], cp.cos(LG_angle + cp.pi)))
        
        refl_event.record()
        refl_event.synchronize()     
        
        
        zd_d = cp.add(z_d, zmirror_d)       
        refl_field_d = self.get_interp_field_GPU(z_d, zd_d, refl_field)
        
        interp_event.record()
        interp_event.synchronize()
        
        
        self.incident_field = cp.asnumpy(inc_field)
        self.zmirror = cp.asnumpy(zmirror_d)
        self.full_envelope = cp.asnumpy(full_envelope)       
        self.refl_field = cp.asnumpy(refl_field_d)
        
        copy_event.record()
        copy_event.synchronize()
        end = datetime.now()
        
        self.timings = [cp.cuda.get_elapsed_time(start_event, copy_in_event), cp.cuda.get_elapsed_time(copy_in_event, inc_event),
                        cp.cuda.get_elapsed_time(inc_event, refl_event),cp.cuda.get_elapsed_time(refl_event, interp_event),
                        cp.cuda.get_elapsed_time(interp_event, copy_event), cp.cuda.get_elapsed_time(start_event, copy_event),
                        (end - start).total_seconds() * 1000]
             
        self.fitted = True
        
        return 

    
    @staticmethod
    def get_interp_field_GPU(z, zd, refl_field):
        """
        Function to interpolate the field to the observer's frame.

        Parameters
        ----------
        z : Array
            Retarded frame
        zd : Array
            Detector frame
        refl_field : Array
            Field

        Returns
        -------
        refl_field_interp : Array
            Interpolated field, sutable for observer frame

        """
        refl_field_interp = cp.zeros_like(refl_field)
        for i in range(z.shape[0]):
            for j in range(z.shape[1]):
                refl_field_interp[i,j,:] = cp.interp(z[i,j,:], zd[i,j,:], refl_field[i,j,:])
        return refl_field_interp    
     
        
    def calculate_spectrum(self):
        """
        This function calculates the spectrum along time/z axis of the field. 
        After this, the follwing is done:
            fft_res - resulted spectrum
            fft_freq - resulted freq array

        Returns
        -------
        None.

        """
        if not self.fitted:
            raise RuntimeError("The model should be calculated before this function call") 
            
        start_event = cp.cuda.Event()
        end_event = cp.cuda.Event()
        start_event.record()
        
        field, d_rate = cp.asarray(self.refl_field), self.d_sampling_rate
        self.fft_res = cp.asnumpy(cp.fft.fft(field, axis=2))
        self.fft_freq = cp.asnumpy(cp.fft.fftfreq(len(field[1,1,:]), d=d_rate))
        
        end_event.record()
        end_event.synchronize()

        self.timings.append(cp.cuda.get_elapsed_time(start_event, end_event))
        self.spectrated = True
        
        return

    def harmonic_cutter(self, number_of_harmonic, width, degree=10):
        """
        Cuts the spectrum with supergauss window with center in number_of_harmonic 
        and width in laser freq. The data to cut is taken from object fields.

        Args
        ----------
        number_of_harmonic : Int
            Number of harmonic to extract.
        width : Float
            Width of the window in laser freq.
        degree : Float
            Supergauss degree for window function

        Returns
        -------
        final_sp : Array
            3d array of cutted spectrum
        final_cut : Array
            Window for cutting
        final_field : Array
            3d array of cutted field
        """
        
        if not self.spectrated:
            raise RuntimeError("Spectrum should be calculated before this function call") 
        
        spectrum = self.fft_res
        degree = 10
        freqs_ = 2*np.pi*self.fft_freq
        width_ = width
        cut1 = np.exp(-(freqs_ - number_of_harmonic)**degree / (width_**degree)) 
        cut2 = np.exp(-(freqs_ + number_of_harmonic)**degree / (width_**degree)) 
        final_cut = cut1 + cut2
        final_sp = final_cut * spectrum
        final_field = np.fft.ifft(final_sp)
        
        return final_sp, final_cut, final_field
        
    def oam_spectrum(self, field, rho_pos=None, z_pos=None):
        
        """
        
        Creates spectrum of the field w.r.t. the azimuthal angle phi. This allow to understand the impact 
        of every OAM harmonic into field. The idea: We are taking the transversal plane at z_pos index 
        and at rho_pos index radius calculate FFT.
        
        Args:
            field - 3d array of field to calculate phi spectrum. Median coordinates for rho and z will be chosen.
            rho_pos and z_pos - indexes of rho and z arrays. 
            
        Returns:
            spectrum - 1d array
            frequency array
        
        """
        if rho_pos is None:
            rho_pos = self.nrho // 2
            print(f"Rho position index was specified as {rho_pos}")
        if z_pos is None:
            z_pos = self.nz // 2
            print(f"Z position index was specified as {z_pos}")
            
        try:        
            phi = self.phi[rho_pos, :, z_pos]
            phi_field = field[rho_pos, :, z_pos]
            
        except IndexError as e:
            print(e)
            print("z_pos or rho_pos index are out of range")
        
        else:
              
            phi_sp = np.fft.fft(phi_field, n = (1000 + self.nphi))
            phi_freq = np.fft.fftfreq((1000 + self.nphi), phi[1] - phi[0])
      
            return phi_sp, phi_freq
        
        
    def field_to_vts(self, harm_field, filename="VTK_file"):
        
        
        """
        
        Function to create VTK file for visualization of the field.
        
        Args:
            harm_field - 3d array of the field.
            filename - name of the file to save.
            
        Returns:
            None, but saves the file.
        
        """
        try:
            
            from pyevtk.hl import pointsToVTK
            from pyevtk.hl import gridToVTK
            
        except:
            print(f"Import failed, install pyevtk.hl")
            

        
        Rho = self.rho[:,:,:]/2/np.pi
        Phi = self.phi[:,:,:]
        Z = self.z[:,:,:]/2/np.pi

        X = Rho * np.cos(Phi)
        Y = Rho * np.sin(Phi)
        x, y, z = X, Y, Z
        
        data = np.real(harm_field)
        data = np.ascontiguousarray(data)
        
        gridToVTK(filename, x/np.amax(x), y/np.amax(y), z/np.amax(z), pointData = {"field_x" : data})
        print('done')
        
        return