vectorwaves 1.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Mayank Soni
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,114 @@
1
+ Metadata-Version: 2.4
2
+ Name: vectorwaves
3
+ Version: 1.0.0
4
+ Summary: Three-dimensional electromagnetic field simulation and topology analysis
5
+ Author: Mayank Soni
6
+ License: MIT License
7
+
8
+ Copyright (c) 2026 Mayank Soni
9
+
10
+ Permission is hereby granted, free of charge, to any person obtaining a copy
11
+ of this software and associated documentation files (the "Software"), to deal
12
+ in the Software without restriction, including without limitation the rights
13
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
14
+ copies of the Software, and to permit persons to whom the Software is
15
+ furnished to do so, subject to the following conditions:
16
+
17
+ The above copyright notice and this permission notice shall be included in all
18
+ copies or substantial portions of the Software.
19
+
20
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
21
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
22
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
23
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
24
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
25
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
26
+ SOFTWARE.
27
+ Project-URL: Homepage, https://github.com/1Rayokelvin/VectorWaves
28
+ Project-URL: Repository, https://github.com/1Rayokelvin/VectorWaves
29
+ Project-URL: Issues, https://github.com/1Rayokelvin/VectorWaves/issues
30
+ Project-URL: Documentation, https://1rayokelvin.github.io/VectorWaves
31
+ Keywords: physics,optics,wave-optics,vector-optics,electromagnetism,electromagnetic-fields,structured-light,gaussian-beams,laguerre-gaussian,plane-wave-expansion,computational-physics,optical-singularities,polarization-singularities,phase-singularities,topological-optics,optical-vortices
32
+ Classifier: Development Status :: 4 - Beta
33
+ Classifier: Intended Audience :: Science/Research
34
+ Classifier: Intended Audience :: Education
35
+ Classifier: License :: OSI Approved :: MIT License
36
+ Classifier: Programming Language :: Python :: 3
37
+ Classifier: Programming Language :: Python :: 3.10
38
+ Classifier: Programming Language :: Python :: 3.11
39
+ Classifier: Programming Language :: Python :: 3.12
40
+ Classifier: Programming Language :: Python :: 3.13
41
+ Classifier: Operating System :: OS Independent
42
+ Classifier: Topic :: Scientific/Engineering :: Physics
43
+ Classifier: Topic :: Scientific/Engineering
44
+ Requires-Python: >=3.10
45
+ Description-Content-Type: text/markdown
46
+ License-File: LICENSE
47
+ Requires-Dist: numpy>=1.24.0
48
+ Requires-Dist: scipy>=1.10.0
49
+ Requires-Dist: numba>=0.59.0
50
+ Provides-Extra: viz
51
+ Requires-Dist: matplotlib>=3.7.0; extra == "viz"
52
+ Requires-Dist: pyvista>=0.40.0; extra == "viz"
53
+ Provides-Extra: progress
54
+ Requires-Dist: tqdm>=4.65.0; extra == "progress"
55
+ Provides-Extra: gpu
56
+ Requires-Dist: cupy-cuda12x>=12.0.0; extra == "gpu"
57
+ Provides-Extra: all
58
+ Requires-Dist: matplotlib>=3.7.0; extra == "all"
59
+ Requires-Dist: pyvista>=0.40.0; extra == "all"
60
+ Requires-Dist: tqdm>=4.65.0; extra == "all"
61
+ Requires-Dist: cupy-cuda12x>=12.0.0; extra == "all"
62
+ Dynamic: license-file
63
+
64
+ VectorWaves provides a framework for generating, computing, and analyzing fully three-dimensional electromagnetic fields and their topology through discrete plane-wave expansions.
65
+
66
+ ## Installation
67
+
68
+ ```bash
69
+ pip install vectorwaves
70
+ ```
71
+
72
+ For additional features, you can install the optional dependencies:
73
+
74
+ | Extra | Purpose |
75
+ |---------|---------|
76
+ | `viz` | Matplotlib and PyVista for visualizations |
77
+ | `progress` | Progress bars via tqdm |
78
+ | `gpu` | CUDA acceleration via CuPy |
79
+ | `all` | All the above |
80
+
81
+ To install,
82
+ ```bash
83
+ pip install vectorwaves[extra]
84
+ ```
85
+
86
+ ## Features
87
+
88
+ - Physics-oriented configuration system with `numpy`, `numba`, and CuPy (GPU) backends.
89
+ - Exact non-paraxial 3D propagation via Fibonacci-sphere discrete plane-wave expansions.
90
+ - Monochromatic and polychromatic sources with arbitrary envelopes, structured light support.
91
+ - Stochastic process generation for speckle-like fields.
92
+ - Fully analytic computation of E-fields, B-fields, spatial derivatives.
93
+ - Topological polarization analysis: C, Cᵀ, and Lᵀ point finding with 3D line tracing.
94
+
95
+ ## Quick Example
96
+
97
+ ```python
98
+ import vectorwaves as vw
99
+
100
+ # Generate a tightly-focused Laguerre-Gaussian beam
101
+ config = vw.get_config()
102
+ config.source.k_space.laguerre_gauss(p=1, l=2, sigma_k_perp=1)
103
+ config.source.randomize.off()
104
+
105
+ # Construct the beam and visualize its plane-wave modes
106
+ beam = vw.setup_beam(config)
107
+
108
+ # Requires matplotlib, install with 'viz' extra: pip install vectorwaves[viz]
109
+ beam.plot_kspace_3d(plot_type='colored_vectors')
110
+ ```
111
+
112
+ ![LG beam k-space](https://github.com/1Rayokelvin/VectorWaves/blob/main/docs/images/LG_kspace.png?raw=true)
113
+
114
+ For tutorials and examples, please refer to the [official documentation](https://1rayokelvin.github.io/VectorWaves). Source code is available on [GitHub](https://github.com/1Rayokelvin/VectorWaves/).
@@ -0,0 +1,83 @@
1
+ # VectorWaves
2
+
3
+ A Python library for constructing and analyzing electromagnetic fields through discrete plane-wave expansions.
4
+
5
+ Classical light is fundamentally an electromagnetic wave. In vacuum, electromagnetic fields admit a plane-wave decomposition, and VectorWaves provides a framework for constructing, computing, and analyzing fully three-dimensional vector fields and their topological structures.
6
+
7
+ **For full documentation, basic usage, and tutorials, visit the [official documentation site](https://1rayokelvin.github.io/VectorWaves).**
8
+
9
+ ## Installation
10
+
11
+ ```bash
12
+ pip install VectorWaves
13
+ ```
14
+
15
+ For additional features, you can install optional dependencies:
16
+
17
+ | Extra | Purpose |
18
+ |---------|---------|
19
+ | `viz` | Matplotlib and PyVista for visualizations |
20
+ | `progress` | Progress bars via tqdm |
21
+ | `gpu` | CUDA acceleration via CuPy |
22
+ | `all` | All the above |
23
+
24
+ To install,
25
+ ```bash
26
+ pip install vectorwaves[#Extra]
27
+ ```
28
+
29
+ ## Features
30
+
31
+ - **Exact 3D Fields**: Electric, Magnetic fields and spatial derivatives via Fibonacci-sphere discrete plane-wave expansions.
32
+ - **Monochromatic Source Models**: Built-in source definitions and arbitrary user-defined spatial spectra.
33
+ - **Polychromatic Source Models**: Built-in spectral distributions and arbitrary user-defined spectral profiles.
34
+ - **Advanced Topology**: Detect and trace polarization singularities (C, Cᵀ, and Lᵀ points/lines).
35
+ - **Physical Diagnostics**: Easily compute Stokes parameters and polarization ellipses for quick visuals.
36
+ - **Stochastic Fields**: Built-in randomization for generating speckles and unpolarized fields.
37
+ - **High Performance**: Hierarchical configuration with seamless switching between CPU (`numpy`, `numba`) and GPU (`cupy`) backends.
38
+
39
+ ## Quick Tour: The Power of Plane-Wave Superposition
40
+
41
+ Because VectorWaves natively superposes true 3D plane-wave vectors, it effortlessly captures complex non-paraxial vector effects that scalar FFT propagators miss entirely.
42
+
43
+ For example, when you tightly focus a linearly polarized Gaussian beam, a longitudinal (z-direction) electric field and a clover-like cross-polarized field naturally emerge purely from the geometry of the converging wavevectors.
44
+
45
+ ```python
46
+ import matplotlib.pyplot as plt
47
+ plt.style.use('dark_background')
48
+ import numpy as np
49
+ import vectorwaves as vw
50
+
51
+ # 1. Configure a tightly focused, x-polarized Gaussian beam
52
+ config = vw.get_config()
53
+ config.op.size = (2,2); config.op.spacing = 0.02
54
+ config.source.pol_vect = (1, 0)
55
+ config.source.k_space.gaussian(sigma_k_perp=2.0) # High divergence
56
+ config.source.theta_max = np.pi/2 # Allow modes up to 90 degrees
57
+ config.source.randomize.off()
58
+
59
+ # 2. Construct the engine and compute the exact vector field at the focus
60
+ engine = vw.setup_engine(config)
61
+ result = engine.compute_on_op(z=0.0)
62
+
63
+ # 3. Extract the intensity of x and z components
64
+ Ex, Ey, Ez = result.E
65
+ I_x, I_z = np.abs(Ex)**2, np.abs(Ez)**2
66
+
67
+ # Plotting the two intensities
68
+ fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12,5))
69
+ extent = engine.op_extent
70
+
71
+ im1 = ax1.imshow(I_x, extent=extent, cmap='magma'); ax1.set_title("|Ex|² (Main)")
72
+ im2 = ax2.imshow(I_z, extent=extent, cmap='magma'); ax2.set_title("|Ez|² (Longitudinal)")
73
+
74
+ fig.colorbar(im1, ax=ax1); fig.colorbar(im2, ax=ax2)
75
+ plt.tight_layout()
76
+ plt.show()
77
+ ```
78
+
79
+ ![Gaussian Beam components](docs/images/gaussian_beam_vector.png)
80
+ *(Note the colorbar scales: the solver automatically calculates the correct relative magnitudes of the emergent longitudinal components without any paraxial approximations).*
81
+
82
+ ### Ready to learn more?
83
+ For a gentler introduction covering basic configurations, spectral envelopes, or fully 3D field visualizations, head over to the **[Getting Started Guide](https://1rayokelvin.github.io/VectorWaves)**.
@@ -0,0 +1,51 @@
1
+ VectorWaves provides a framework for generating, computing, and analyzing fully three-dimensional electromagnetic fields and their topology through discrete plane-wave expansions.
2
+
3
+ ## Installation
4
+
5
+ ```bash
6
+ pip install vectorwaves
7
+ ```
8
+
9
+ For additional features, you can install the optional dependencies:
10
+
11
+ | Extra | Purpose |
12
+ |---------|---------|
13
+ | `viz` | Matplotlib and PyVista for visualizations |
14
+ | `progress` | Progress bars via tqdm |
15
+ | `gpu` | CUDA acceleration via CuPy |
16
+ | `all` | All the above |
17
+
18
+ To install,
19
+ ```bash
20
+ pip install vectorwaves[extra]
21
+ ```
22
+
23
+ ## Features
24
+
25
+ - Physics-oriented configuration system with `numpy`, `numba`, and CuPy (GPU) backends.
26
+ - Exact non-paraxial 3D propagation via Fibonacci-sphere discrete plane-wave expansions.
27
+ - Monochromatic and polychromatic sources with arbitrary envelopes, structured light support.
28
+ - Stochastic process generation for speckle-like fields.
29
+ - Fully analytic computation of E-fields, B-fields, spatial derivatives.
30
+ - Topological polarization analysis: C, Cᵀ, and Lᵀ point finding with 3D line tracing.
31
+
32
+ ## Quick Example
33
+
34
+ ```python
35
+ import vectorwaves as vw
36
+
37
+ # Generate a tightly-focused Laguerre-Gaussian beam
38
+ config = vw.get_config()
39
+ config.source.k_space.laguerre_gauss(p=1, l=2, sigma_k_perp=1)
40
+ config.source.randomize.off()
41
+
42
+ # Construct the beam and visualize its plane-wave modes
43
+ beam = vw.setup_beam(config)
44
+
45
+ # Requires matplotlib, install with 'viz' extra: pip install vectorwaves[viz]
46
+ beam.plot_kspace_3d(plot_type='colored_vectors')
47
+ ```
48
+
49
+ ![LG beam k-space](https://github.com/1Rayokelvin/VectorWaves/blob/main/docs/images/LG_kspace.png?raw=true)
50
+
51
+ For tutorials and examples, please refer to the [official documentation](https://1rayokelvin.github.io/VectorWaves). Source code is available on [GitHub](https://github.com/1Rayokelvin/VectorWaves/).
@@ -0,0 +1,101 @@
1
+ [build-system]
2
+ requires = ["setuptools>=68"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "vectorwaves"
7
+ dynamic = ["version"]
8
+ description = "Three-dimensional electromagnetic field simulation and topology analysis"
9
+ authors = [
10
+ { name = "Mayank Soni" }
11
+ ]
12
+ readme = "README_PyPi.md"
13
+ license = { file = "LICENSE" }
14
+ requires-python = ">=3.10"
15
+
16
+ # CORE DEPENDENCIES
17
+ # These are strictly required for the physics pipeline to function.
18
+ dependencies = [
19
+ "numpy>=1.24.0",
20
+ "scipy>=1.10.0",
21
+ "numba>=0.59.0"
22
+ ]
23
+
24
+ classifiers = [
25
+ "Development Status :: 4 - Beta",
26
+
27
+ "Intended Audience :: Science/Research",
28
+ "Intended Audience :: Education",
29
+
30
+ "License :: OSI Approved :: MIT License",
31
+
32
+ "Programming Language :: Python :: 3",
33
+ "Programming Language :: Python :: 3.10",
34
+ "Programming Language :: Python :: 3.11",
35
+ "Programming Language :: Python :: 3.12",
36
+ "Programming Language :: Python :: 3.13",
37
+
38
+ "Operating System :: OS Independent",
39
+
40
+ "Topic :: Scientific/Engineering :: Physics",
41
+ "Topic :: Scientific/Engineering",
42
+ ]
43
+
44
+ keywords = [
45
+ "physics",
46
+ "optics",
47
+ "wave-optics",
48
+ "vector-optics",
49
+ "electromagnetism",
50
+ "electromagnetic-fields",
51
+ "structured-light",
52
+ "gaussian-beams",
53
+ "laguerre-gaussian",
54
+ "plane-wave-expansion",
55
+ "computational-physics",
56
+ "optical-singularities",
57
+ "polarization-singularities",
58
+ "phase-singularities",
59
+ "topological-optics",
60
+ "optical-vortices"
61
+ ]
62
+
63
+ [project.optional-dependencies]
64
+ # VISUALIZATION: Plotting k-space, profiles, and tutorials
65
+ viz = [
66
+ "matplotlib>=3.7.0",
67
+ "pyvista>=0.40.0"
68
+ ]
69
+
70
+ # PROGRESS: Console loading bars
71
+ progress = [
72
+ "tqdm>=4.65.0"
73
+ ]
74
+
75
+ # GPU: CuPy for CUDA 12.x (Note: Users with AMD or older CUDA will need to install their specific cupy build)
76
+ gpu = [
77
+ "cupy-cuda12x>=12.0.0"
78
+ ]
79
+
80
+ # ALL optional dependencies
81
+ all = [
82
+ "matplotlib>=3.7.0",
83
+ "pyvista>=0.40.0",
84
+ "tqdm>=4.65.0",
85
+ "cupy-cuda12x>=12.0.0"
86
+ ]
87
+
88
+ [tool.setuptools.dynamic]
89
+ version = { attr = "vectorwaves.version.current_version_str" }
90
+
91
+ [tool.setuptools]
92
+ package-dir = { "" = "src" }
93
+
94
+ [tool.setuptools.packages.find]
95
+ where = ["src"]
96
+
97
+ [project.urls]
98
+ Homepage = "https://github.com/1Rayokelvin/VectorWaves"
99
+ Repository = "https://github.com/1Rayokelvin/VectorWaves"
100
+ Issues = "https://github.com/1Rayokelvin/VectorWaves/issues"
101
+ Documentation = "https://1rayokelvin.github.io/VectorWaves"
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,88 @@
1
+ """
2
+ VectorWaves - A Python library for electromagnetic field simulations using plane wave expansion.
3
+ ==========================================================================
4
+
5
+ Main workflow:
6
+ >>> import vectorwaves as vw
7
+ >>> config = vw.get_config()
8
+ >>> config.source.num_modes = 10000
9
+ >>> expt = vw.setup_engine(config)
10
+ >>> result = expt.compute_on_op(z=0.0)
11
+
12
+ Author: Mayank Soni
13
+ Year: 2026
14
+ """
15
+ from .version import current_version, current_version_str
16
+ __version_info__ = current_version
17
+ __version__ = current_version_str
18
+
19
+ # Core configuration
20
+ from .config_stuff import (
21
+ Config,
22
+ OpConfig,
23
+ SourceConfig,
24
+ RandomizeConfig,
25
+ KSpaceConfig,
26
+ PolychromaticConfig,
27
+ get_config,
28
+ load_config
29
+ )
30
+
31
+ # Core Pipeline
32
+ from .beam_stuff import BeamMaker, Beam
33
+ from .engine_stuff import FieldEngine, FieldResult, tqdm
34
+
35
+ # Utilities & Physics Math
36
+ from .utils import (
37
+ get_stokes_params,
38
+ get_pol_ellipse_params,
39
+ decompose_in_basis,
40
+ )
41
+
42
+ # Singularity Finders
43
+ from .singularities import SingularityFinder
44
+
45
+ def setup_beam(config: Config) -> Beam:
46
+ """High-level wrapper to generate the beam from config object."""
47
+ return BeamMaker(config).generate_beam()
48
+
49
+ def setup_engine(config: Config) -> FieldEngine:
50
+ """
51
+ High-level wrapper to generate the beam and initialize the computation engine.
52
+ """
53
+ return FieldEngine(setup_beam(config), config)
54
+
55
+ __all__ =[
56
+ # Config
57
+ "Config",
58
+ "OpConfig",
59
+ "SourceConfig",
60
+ "RandomizeConfig",
61
+ "KSpaceConfig",
62
+ "PolychromaticConfig",
63
+ "get_config",
64
+ "load_config",
65
+
66
+ # Engine Pipeline
67
+ "setup_engine",
68
+ "setup_beam",
69
+ "BeamMaker",
70
+ "Beam",
71
+ "FieldEngine",
72
+ "FieldResult",
73
+
74
+ # Utilities
75
+ "get_stokes_params",
76
+ "get_pol_ellipse_params",
77
+ "decompose_in_basis",
78
+
79
+ # Topologies
80
+ "SingularityFinder",
81
+
82
+ # tqdm
83
+ "tqdm",
84
+
85
+ # Version
86
+ "__version__",
87
+ "__version_info__"
88
+ ]
File without changes
@@ -0,0 +1,197 @@
1
+ import numpy as np
2
+ try:
3
+ import cupy as cp
4
+ has_cupy = True
5
+ except ImportError:
6
+ has_cupy = False
7
+
8
+ class CupyMethods:
9
+ def __init__(self, beam, use_single_precision=True):
10
+ if not has_cupy:
11
+ raise RuntimeError("CuPy/CUDA not found.")
12
+
13
+ self.use_single = use_single_precision
14
+ self.real_dt = cp.float32 if self.use_single else cp.float64
15
+ self.comp_dt = cp.complex64 if self.use_single else cp.complex128
16
+
17
+ # --- Persistent Model Data ---
18
+ def to_gpu(arr, dtype): return cp.ascontiguousarray(cp.asarray(arr, dtype=dtype))
19
+
20
+ self.kx = to_gpu(beam.k[0], self.real_dt)
21
+ self.ky = to_gpu(beam.k[1], self.real_dt)
22
+ self.kz = to_gpu(beam.k[2], self.real_dt)
23
+ self.w = to_gpu(beam.w, self.real_dt)
24
+ self.inv_w = to_gpu(beam.inv_w, self.real_dt)
25
+ self.c_base = to_gpu(beam.c, self.comp_dt)
26
+ self.num_waves = len(self.w)
27
+
28
+ self._kernel_cache = {}
29
+
30
+ def _get_kernel(self, num_components, is_grid=False):
31
+ key = (num_components, is_grid)
32
+ if key in self._kernel_cache: return self._kernel_cache[key]
33
+
34
+ real_t = "float" if self.use_single else "double"
35
+ comp_t = "complex<float>" if self.use_single else "complex<double>"
36
+ sincos_f = "sincosf" if self.use_single else "sincos"
37
+
38
+ coord_logic = """
39
+ int ix = p % nx;
40
+ int iy = p / nx;
41
+ {real_t} px = x_vec[ix];
42
+ {real_t} py = y_vec[iy];
43
+ {real_t} pz = z_scalar;
44
+ """ if is_grid else """
45
+ {real_t} px = x_vec[p];
46
+ {real_t} py = y_vec[p];
47
+ {real_t} pz = z_vec[p];
48
+ """
49
+
50
+ kernel_code = f"""
51
+ #include <cupy/complex.cuh>
52
+ extern "C" __global__
53
+ void compute_kernel(
54
+ const {real_t}* __restrict__ x_vec, const {real_t}* __restrict__ y_vec, const {real_t}* __restrict__ z_vec,
55
+ const {real_t}* __restrict__ kx, const {real_t}* __restrict__ ky, const {real_t}* __restrict__ kz,
56
+ const {real_t}* __restrict__ w, const {comp_t}* __restrict__ super_vec,
57
+ {comp_t}* __restrict__ out,
58
+ {real_t} z_scalar, {real_t} t, int nx, int num_pts, int num_waves
59
+ ) {{
60
+ int p = blockDim.x * blockIdx.x + threadIdx.x;
61
+ if (p >= num_pts) return;
62
+
63
+ {coord_logic.format(real_t=real_t)}
64
+
65
+ {comp_t} acc[{num_components}];
66
+ #pragma unroll
67
+ for(int i=0; i<{num_components}; i++) acc[i] = {comp_t}(0, 0);
68
+
69
+ for (int i = 0; i < num_waves; i++) {{
70
+ {real_t} phase = kx[i]*px + ky[i]*py + kz[i]*pz - w[i]*t;
71
+ {real_t} s, c;
72
+ {sincos_f}(phase, &s, &c);
73
+ {comp_t} wf(c, s);
74
+
75
+ #pragma unroll
76
+ for (int j = 0; j < {num_components}; j++) {{
77
+ acc[j] += super_vec[i * {num_components} + j] * wf;
78
+ }}
79
+ }}
80
+
81
+ for (int j = 0; j < {num_components}; j++) {{
82
+ out[j * num_pts + p] = acc[j];
83
+ }}
84
+ }}
85
+ """
86
+ kernel = cp.RawKernel(kernel_code, 'compute_kernel')
87
+ self._kernel_cache[key] = kernel
88
+ return kernel
89
+
90
+ def _prepare_super_vec(self, need_b, need_derivs):
91
+ vecs = [self.c_base]
92
+ if need_b:
93
+ bx = (self.ky * self.c_base[2] - self.kz * self.c_base[1]) * self.inv_w
94
+ by = (self.kz * self.c_base[0] - self.kx * self.c_base[2]) * self.inv_w
95
+ bz = (self.kx * self.c_base[1] - self.ky * self.c_base[0]) * self.inv_w
96
+ vecs.append(cp.stack([bx, by, bz]))
97
+ if need_derivs:
98
+ ik = 1j * cp.stack([self.kx, self.ky, self.kz])
99
+ for i in range(3): vecs.append(self.c_base * ik[i])
100
+
101
+ return cp.ascontiguousarray(cp.vstack(vecs).T)
102
+
103
+ def compute_grid(self, x_vec, y_vec, z, t, need_b=True, need_derivs=True, progress_callback=None):
104
+ nx, ny = len(x_vec), len(y_vec)
105
+
106
+ super_vec = self._prepare_super_vec(need_b, need_derivs)
107
+ num_comps = super_vec.shape[1]
108
+ kernel = self._get_kernel(num_comps, is_grid=True)
109
+
110
+ # Pre-allocate the final CPU array to store results
111
+ out_h = np.zeros((num_comps, ny, nx), dtype=np.complex128)
112
+
113
+ # Move the X vector to GPU once
114
+ x_g = cp.asarray(x_vec, dtype=self.real_dt)
115
+
116
+ # Target ~500MB maximum VRAM footprint for the output buffer
117
+ bytes_per_element = 8 if self.use_single else 16
118
+ bytes_per_row = nx * num_comps * bytes_per_element
119
+ MAX_VRAM_BYTES = 1000 * 1024 * 1024 # 1 GB
120
+
121
+ # Calculate how many rows we can safely process at once
122
+ rows_per_batch = max(1, MAX_VRAM_BYTES // bytes_per_row)
123
+
124
+ for i in range(0, ny, rows_per_batch):
125
+ end = min(i + rows_per_batch, ny)
126
+ cur_ny = end - i
127
+ cur_pts = nx * cur_ny
128
+
129
+ # Move just this chunk of Y coordinates to GPU
130
+ y_g = cp.asarray(y_vec[i:end], dtype=self.real_dt)
131
+
132
+ # Allocate GPU buffer for just this batch
133
+ out_g = cp.empty((num_comps, cur_pts), dtype=self.comp_dt)
134
+
135
+ threads = 256
136
+ blocks = (cur_pts + threads - 1) // threads
137
+
138
+ kernel((blocks,), (threads,), (
139
+ x_g, y_g, None, self.kx, self.ky, self.kz, self.w,
140
+ super_vec, out_g,
141
+ self.real_dt(z), self.real_dt(t), cp.int32(nx), cp.int32(cur_pts), cp.int32(self.num_waves)
142
+ ))
143
+
144
+ # Fetch result to CPU and reshape directly into the pre-allocated CPU array
145
+ out_h[:, i:end, :] = out_g.get().reshape(num_comps, cur_ny, nx)
146
+
147
+ if progress_callback:
148
+ progress_callback(cur_ny)
149
+
150
+ # Unpack the CPU array into standard shapes
151
+ E = out_h[0:3]
152
+ idx = 3
153
+ B = out_h[idx:idx+3] if need_b else None
154
+ if need_b: idx += 3
155
+ D = (out_h[idx:idx+3], out_h[idx+3:idx+6], out_h[idx+6:idx+9]) if need_derivs else (None,None,None)
156
+
157
+ return E, D, B
158
+
159
+ def compute_cloud(self, x, y, z, t, need_b=True, need_derivs=True, progress_callback=None):
160
+ num_pts = len(x)
161
+ super_vec = self._prepare_super_vec(need_b, need_derivs)
162
+ num_comps = super_vec.shape[1]
163
+ kernel = self._get_kernel(num_comps, is_grid=False)
164
+
165
+ E_h, B_h, dx_h, dy_h, dz_h = self._allocate_cpu_arrays((num_pts,))
166
+ CHUNK = 500_000
167
+
168
+ for s in range(0, num_pts, CHUNK):
169
+ e = min(s + CHUNK, num_pts)
170
+ cur_n = e - s
171
+ x_g, y_g, z_g = [cp.asarray(arr[s:e], dtype=self.real_dt) for arr in [x,y,z]]
172
+ out_g = cp.empty((num_comps, cur_n), dtype=self.comp_dt)
173
+
174
+ kernel(((cur_n+255)//256,), (256,), (
175
+ x_g, y_g, z_g, self.kx, self.ky, self.kz, self.w,
176
+ super_vec, out_g, 0.0, self.real_dt(t), 0, cp.int32(cur_n), cp.int32(self.num_waves)
177
+ ))
178
+
179
+ out_c = out_g.get()
180
+ if progress_callback: progress_callback(cur_n)
181
+ E_h[:, s:e] = out_c[0:3]
182
+ idx = 3
183
+ if need_b: B_h[:, s:e] = out_c[idx:idx+3]; idx += 3
184
+ if need_derivs:
185
+ dx_h[:, s:e], dy_h[:, s:e], dz_h[:, s:e] = out_c[idx:idx+3], out_c[idx+3:idx+6], out_c[idx+6:idx+9]
186
+
187
+ return E_h, (dx_h, dy_h, dz_h) if need_derivs else (None,None,None), B_h if need_b else None
188
+
189
+ def _allocate_cpu_arrays(self, shape):
190
+ return tuple(np.zeros((3, *shape), dtype=np.complex128) for _ in range(5))
191
+
192
+ def compute_point(self, x, y, z, t, need_b=True, need_derivs=True):
193
+ E, D, B = self.compute_cloud(np.array([x]), np.array([y]), np.array([z]), t, need_b, need_derivs)
194
+ return E[:,0], (tuple(d[:,0] for d in D) if need_derivs else (None,None,None)), (B[:,0] if need_b else None) # type: ignore
195
+
196
+ def __del__(self):
197
+ self._kernel_cache.clear()