LG_generation.py - dev - stronius/OAM_ROM_Model

LG_generation.py
534 строки18 KB
stronius
minor bugs fix: degree and extra interp func
23 сен 2025, 13:41
23 сен 2025, 13:41f2845f9
from datetime import datetime
import numpy as np
from numba import jit, float64, complex128, int64, types
from scipy.special import genlaguerre

from dataclasses import dataclass
from math import factorial

from typing import List, Tuple, Literal

# test

@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:
        BeamConfig: 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

@jit(types.Tuple((complex128[:,:,:], float64[:,:,:], int64, int64))
     (float64[:,:,:], 
      float64[:,:,:], 
      float64[:,:,:],
      float64,
      int64,
      int64,
      float64,
      float64,
      float64), 
     cache=True, 
     nopython=True, 
     parallel=True)
def LG_mode_accelerated(r, phi, z, mode_params_a0, mode_params_m, 
                        mode_params_p, 
                        mode_params_Mp, 
                        mode_params_w0, 
                        tau):
    """
    Calculates LG mode field with numba 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 = np.square(r)
    z_sq = np.square(z)
    
    w = w0 * np.sqrt(1 + z_sq / Z_R**2)
    w_sq = np.square(w)
    R_z = z * (1 + Z_R**2 / z_sq)
    chi = (2 * p + np.abs(m) + 1) * np.arctan(z / Z_R)
    
    toret = (w0 / w) * (r * np.sqrt(2) / w)**(np.abs(m)) * np.exp(-r_sq / w_sq)
    toret = toret * np.exp(-1j * Mp * z + 1j * m * phi + 1j * r_sq / 2 / R_z - 1j*chi - z_sq / 2 / tau**2)

    return a0*toret, 2*r_sq/w_sq, p, m        
        
class LG_mode_generation():
    
    """
    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(r, phi, z, mode_params):
        mode, arg, p,m = LG_mode_accelerated(r, phi, z, *mode_params)

        Lpm = genlaguerre(p, np.abs(m))
        coeff = np.sqrt(2 * factorial(p) / (np.pi * factorial(abs(m) + p)))
        mode = mode * Lpm(arg) * coeff
        return mode
           
        
    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. 
        """
        
        r, phi, z = self.cartesian_to_polar()
        start = datetime.now()
        print("Calculating LG modes...", end="")
        
        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(r, phi, z-delay, mode)

            LG_abs = np.abs(LG)
            LG_angle = np.angle(LG)
            
            inc_field += LG_abs * np.cos(LG_angle)
            cos_sum += inc_field
            sin_sum += LG_abs * np.sin(LG_angle)
            
            LG_envs[i] = LG_abs
                    
        full_envelope = np.sqrt(cos_sum**2 + sin_sum**2)
        
        inc_time = datetime.now()
        print(f"Incident field calculated in {(inc_time - start).total_seconds()} seconds")
        
        zmirror = cos_sum**2 / (1 + full_envelope**2)
        
        for i, (mode, delay) in enumerate(zip(modes, delays)):
            
            LG = self.LG_mode(r, phi, z-delay - zmirror, mode)

            LG_angle = np.angle(LG)
           
            refl_field += LG_envs[i] * np.cos(LG_angle + np.pi)
        
        refl_time = datetime.now()
        print(f"Refl field calculated in {(refl_time - inc_time).total_seconds()} seconds")
        
        zd = self.z + zmirror
        
        self.incident_field = inc_field
        self.zmirror = zmirror
        self.full_envelope = full_envelope
        
        self.refl_field = self.get_interp_field_accelerated(self.z, zd, refl_field)
        
        print(f"Interp field calculated in {(datetime.now() - refl_time).total_seconds()} seconds")
        
        self.fitted = True
        print(f"OK!\nFinished training in {(datetime.now()-start).total_seconds()} seconds.")
        
        return 

    @staticmethod
    @jit(complex128[:,:,:](float64[:,:,:], float64[:,:,:], float64[:,:,:]), 
         nopython=True, 
         cache=True)
    def get_interp_field_accelerated(z, zd, refl_field):
        """
        Speed-up version of interpolation function with Numba.
    
        Args:
        ----------
        z : float64[:,:,:]
            Coordinates retarded frame.
        zd : float64[:,:,:]
            Coordinates detector frame.
        refl_field : complex128[:,:,:]
            Reflected field.
    
        Returns:
        -----------
        refl_field_interp : complex128[:,:,:]
            Interpolated field, sutable for observer frame
        """
        nz, nphi, nrho = z.shape
        refl_field_interp = np.zeros_like(refl_field, dtype=np.complex128)
        
        for i in range(nz):
            for j in range(nphi):
                refl_field_interp[i, j, :] = np.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 = datetime.now()
        
        field, d_rate = self.refl_field, self.d_sampling_rate
        self.fft_res = np.fft.fft(field, axis=2)
        self.fft_freq = np.fft.fftfreq(len(field[1,1,:]), d=d_rate)
        end = datetime.now()
        print(f"time = {(end - start).total_seconds()}")
        
        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
        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
        print(final_cut.shape, spectrum.shape)
        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:
            pad = 0
            phi_sp = np.fft.fft(phi_field, n = (pad + self.nphi))
            phi_freq = np.fft.fftfreq((pad + self.nphi), phi[1] - phi[0])
      
            return phi_sp, phi_freq
        
        
    def field_to_vts(self, harm_field, style="3d", 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
            from pyevtk.hl import imageToVTK
            
        except:
            print(f"Import failed, install pyevtk.hl")
            

        if style == "3d":
            
            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
        
        elif style == "2d":
            
            Rho = self.rho[:,:,:]/2/np.pi
            Phi = self.phi[:,:,:]
            Z = np.array([0.0])

            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
        else:
            print("wrong style, use '3d' or '2d'")