1use crate::GpuError;
14use amari_relativistic::{particle::RelativisticParticle, spacetime::SpacetimeVector};
15use bytemuck::{Pod, Zeroable};
16use wgpu::util::DeviceExt;
17
18#[repr(C)]
20#[derive(Copy, Clone, Debug, Pod, Zeroable)]
21pub struct GpuSpacetimeVector {
22 pub t: f32,
24 pub x: f32,
26 pub y: f32,
28 pub z: f32,
30}
31
32impl GpuSpacetimeVector {
33 pub fn new(t: f32, x: f32, y: f32, z: f32) -> Self {
35 Self { t, x, y, z }
36 }
37
38 pub fn from_spacetime_vector(sv: &SpacetimeVector) -> Self {
40 let coords = sv.coordinates();
41 Self::new(
42 coords[0] as f32,
43 coords[1] as f32,
44 coords[2] as f32,
45 coords[3] as f32,
46 )
47 }
48
49 pub fn to_spacetime_vector(&self) -> SpacetimeVector {
51 SpacetimeVector::from_coordinates([
52 self.t as f64,
53 self.x as f64,
54 self.y as f64,
55 self.z as f64,
56 ])
57 }
58
59 fn is_finite(&self) -> bool {
60 self.t.is_finite() && self.x.is_finite() && self.y.is_finite() && self.z.is_finite()
61 }
62}
63
64#[repr(C)]
66#[derive(Copy, Clone, Debug, Pod, Zeroable)]
67pub struct GpuRelativisticParticle {
68 pub position: GpuSpacetimeVector,
70 pub velocity: GpuSpacetimeVector,
72 pub mass: f32,
74 pub charge: f32,
76 pub proper_time: f32,
78 pub _padding: [f32; 5],
80}
81
82#[repr(C)]
84#[derive(Copy, Clone, Debug, Pod, Zeroable)]
85pub struct GpuTrajectoryParams {
86 pub dt: f32,
88 pub steps: u32,
90 pub tolerance: f32,
92 pub renorm_freq: u32,
94 pub schwarzschild_radius: f32,
96 pub gm_parameter: f32,
98 pub _padding: [f32; 2],
100}
101
102pub struct GpuRelativisticPhysics {
104 device: wgpu::Device,
105 queue: wgpu::Queue,
106 spacetime_pipeline: wgpu::ComputePipeline,
107 geodesic_pipeline: wgpu::ComputePipeline,
108 #[allow(dead_code)]
109 trajectory_pipeline: wgpu::ComputePipeline,
110}
111
112impl GpuRelativisticPhysics {
113 pub async fn new() -> Result<Self, GpuError> {
115 let instance = wgpu::Instance::default();
116
117 let adapter = instance
118 .request_adapter(&wgpu::RequestAdapterOptions::default())
119 .await
120 .ok_or_else(|| {
121 GpuError::InitializationError("No suitable GPU adapter found".to_string())
122 })?;
123
124 let (device, queue) = adapter
125 .request_device(
126 &wgpu::DeviceDescriptor {
127 label: Some("Relativistic Physics GPU"),
128 required_features: wgpu::Features::empty(),
129 required_limits: wgpu::Limits::default(),
130 },
131 None,
132 )
133 .await
134 .map_err(|e| {
135 GpuError::InitializationError(format!("Failed to create device: {}", e))
136 })?;
137
138 let spacetime_pipeline = Self::create_spacetime_pipeline(&device)?;
140 let geodesic_pipeline = Self::create_geodesic_pipeline(&device)?;
141 let trajectory_pipeline = Self::create_trajectory_pipeline(&device)?;
142
143 Ok(Self {
144 device,
145 queue,
146 spacetime_pipeline,
147 geodesic_pipeline,
148 trajectory_pipeline,
149 })
150 }
151
152 fn create_spacetime_pipeline(device: &wgpu::Device) -> Result<wgpu::ComputePipeline, GpuError> {
154 let shader_source = r#"
155 @group(0) @binding(0) var<storage, read_write> vectors: array<vec4<f32>>;
156 @group(0) @binding(1) var<storage, read_write> results: array<f32>;
157
158 @compute @workgroup_size(64)
159 fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
160 let index = global_id.x;
161 if (index >= arrayLength(&vectors)) {
162 return;
163 }
164
165 let v = vectors[index];
166
167 // Minkowski inner product: t² - x² - y² - z²
168 let minkowski_norm_sq = v.x * v.x - v.y * v.y - v.z * v.z - v.w * v.w;
169 results[index] = minkowski_norm_sq;
170
171 // Normalize four-velocity if needed (u·u = c²)
172 let c_sq = 299792458.0 * 299792458.0;
173 if (abs(minkowski_norm_sq - c_sq) > 1e-6) {
174 let norm = sqrt(abs(minkowski_norm_sq));
175 if (norm > 1e-12) {
176 let factor = sqrt(c_sq) / norm;
177 vectors[index] = vec4<f32>(v.x * factor, v.y * factor, v.z * factor, v.w * factor);
178 }
179 }
180 }
181 "#;
182
183 let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
184 label: Some("Spacetime Algebra Compute Shader"),
185 source: wgpu::ShaderSource::Wgsl(shader_source.into()),
186 });
187
188 let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
189 label: Some("Spacetime Bind Group Layout"),
190 entries: &[
191 wgpu::BindGroupLayoutEntry {
192 binding: 0,
193 visibility: wgpu::ShaderStages::COMPUTE,
194 ty: wgpu::BindingType::Buffer {
195 ty: wgpu::BufferBindingType::Storage { read_only: false },
196 has_dynamic_offset: false,
197 min_binding_size: None,
198 },
199 count: None,
200 },
201 wgpu::BindGroupLayoutEntry {
202 binding: 1,
203 visibility: wgpu::ShaderStages::COMPUTE,
204 ty: wgpu::BindingType::Buffer {
205 ty: wgpu::BufferBindingType::Storage { read_only: false },
206 has_dynamic_offset: false,
207 min_binding_size: None,
208 },
209 count: None,
210 },
211 ],
212 });
213
214 let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
215 label: Some("Spacetime Pipeline Layout"),
216 bind_group_layouts: &[&bind_group_layout],
217 push_constant_ranges: &[],
218 });
219
220 Ok(
221 device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
222 label: Some("Spacetime Compute Pipeline"),
223 layout: Some(&pipeline_layout),
224 module: &shader,
225 entry_point: "main",
226 }),
227 )
228 }
229
230 fn create_geodesic_pipeline(device: &wgpu::Device) -> Result<wgpu::ComputePipeline, GpuError> {
232 let shader_source = r#"
233 struct Particle {
234 position: vec4<f32>,
235 velocity: vec4<f32>,
236 mass: f32,
237 charge: f32,
238 proper_time: f32,
239 padding: array<f32, 5>,
240 };
241
242 struct TrajectoryParams {
243 dt: f32,
244 steps: u32,
245 tolerance: f32,
246 renorm_freq: u32,
247 rs: f32,
248 gm: f32,
249 padding: vec2<f32>,
250 };
251
252 @group(0) @binding(0) var<storage, read_write> particles: array<Particle>;
253 @group(0) @binding(1) var<uniform> params: TrajectoryParams;
254
255 // Schwarzschild metric Christoffel symbols (simplified)
256 fn christoffel_t_rr(r: f32, rs: f32) -> f32 {
257 let factor = rs / (2.0 * r * r);
258 return factor * (1.0 - rs / r);
259 }
260
261 fn christoffel_r_tr(r: f32, rs: f32) -> f32 {
262 return rs / (2.0 * r * r * (1.0 - rs / r));
263 }
264
265 fn christoffel_r_rr(r: f32, rs: f32) -> f32 {
266 return -rs / (2.0 * r * (r - rs));
267 }
268
269 @compute @workgroup_size(64)
270 fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
271 let index = global_id.x;
272 if (index >= arrayLength(&particles)) {
273 return;
274 }
275
276 var particle = particles[index];
277 let pos = particle.position;
278 let vel = particle.velocity;
279
280 // Compute spatial radius
281 let r = sqrt(pos.y * pos.y + pos.z * pos.z + pos.w * pos.w);
282
283 if (r < params.rs * 1.1) {
284 // Too close to singularity, skip
285 return;
286 }
287
288 // Velocity Verlet step for geodesic equation
289 // Simplified for Schwarzschild metric
290
291 // Compute acceleration components
292 let c_trr = christoffel_t_rr(r, params.rs);
293 let c_rtr = christoffel_r_tr(r, params.rs);
294 let c_rrr = christoffel_r_rr(r, params.rs);
295
296 // Geodesic equation: d²x^μ/dτ² = -Γ^μ_αβ v^α v^β
297 var accel = vec4<f32>(0.0, 0.0, 0.0, 0.0);
298
299 // Simplified acceleration calculation
300 accel.x = -c_trr * vel.y * vel.y; // dt component
301 accel.y = -c_rtr * vel.x * vel.y - c_rrr * vel.y * vel.y; // dr component
302
303 // Update position and velocity
304 let dt = params.dt;
305 particle.position = pos + vel * dt + 0.5 * accel * dt * dt;
306 particle.velocity = vel + accel * dt;
307
308 // Renormalize four-velocity periodically
309 let step_mod = u32(particle.proper_time / dt) % params.renorm_freq;
310 if (step_mod == 0u) {
311 let c_sq = 299792458.0 * 299792458.0;
312 let norm_sq = particle.velocity.x * particle.velocity.x -
313 particle.velocity.y * particle.velocity.y -
314 particle.velocity.z * particle.velocity.z -
315 particle.velocity.w * particle.velocity.w;
316
317 if (abs(norm_sq - c_sq) > params.tolerance) {
318 let norm = sqrt(abs(norm_sq));
319 if (norm > 1e-12) {
320 let factor = sqrt(c_sq) / norm;
321 particle.velocity *= factor;
322 }
323 }
324 }
325
326 particle.proper_time += dt;
327 particles[index] = particle;
328 }
329 "#;
330
331 let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
332 label: Some("Geodesic Integration Compute Shader"),
333 source: wgpu::ShaderSource::Wgsl(shader_source.into()),
334 });
335
336 let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
337 label: Some("Geodesic Bind Group Layout"),
338 entries: &[
339 wgpu::BindGroupLayoutEntry {
340 binding: 0,
341 visibility: wgpu::ShaderStages::COMPUTE,
342 ty: wgpu::BindingType::Buffer {
343 ty: wgpu::BufferBindingType::Storage { read_only: false },
344 has_dynamic_offset: false,
345 min_binding_size: None,
346 },
347 count: None,
348 },
349 wgpu::BindGroupLayoutEntry {
350 binding: 1,
351 visibility: wgpu::ShaderStages::COMPUTE,
352 ty: wgpu::BindingType::Buffer {
353 ty: wgpu::BufferBindingType::Uniform,
354 has_dynamic_offset: false,
355 min_binding_size: None,
356 },
357 count: None,
358 },
359 ],
360 });
361
362 let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
363 label: Some("Geodesic Pipeline Layout"),
364 bind_group_layouts: &[&bind_group_layout],
365 push_constant_ranges: &[],
366 });
367
368 Ok(
369 device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
370 label: Some("Geodesic Compute Pipeline"),
371 layout: Some(&pipeline_layout),
372 module: &shader,
373 entry_point: "main",
374 }),
375 )
376 }
377
378 fn create_trajectory_pipeline(
380 device: &wgpu::Device,
381 ) -> Result<wgpu::ComputePipeline, GpuError> {
382 Self::create_geodesic_pipeline(device)
384 }
385
386 pub async fn compute_minkowski_products(
390 &self,
391 vectors: &[GpuSpacetimeVector],
392 ) -> Result<Vec<f32>, GpuError> {
393 if vectors.is_empty() {
394 return Ok(Vec::new());
395 }
396 for (index, vector) in vectors.iter().enumerate() {
397 if !vector.is_finite() {
398 return Err(GpuError::BufferError(format!(
399 "spacetime vector {index} contains non-finite values"
400 )));
401 }
402 }
403
404 let vectors_buffer = self
405 .device
406 .create_buffer_init(&wgpu::util::BufferInitDescriptor {
407 label: Some("Spacetime Vectors Buffer"),
408 contents: bytemuck::cast_slice(vectors),
409 usage: wgpu::BufferUsages::STORAGE
410 | wgpu::BufferUsages::COPY_DST
411 | wgpu::BufferUsages::COPY_SRC,
412 });
413
414 let results_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
415 label: Some("Results Buffer"),
416 size: (vectors.len() * std::mem::size_of::<f32>()) as u64,
417 usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
418 mapped_at_creation: false,
419 });
420
421 let bind_group_layout = self.spacetime_pipeline.get_bind_group_layout(0);
422 let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
423 label: Some("Spacetime Bind Group"),
424 layout: &bind_group_layout,
425 entries: &[
426 wgpu::BindGroupEntry {
427 binding: 0,
428 resource: vectors_buffer.as_entire_binding(),
429 },
430 wgpu::BindGroupEntry {
431 binding: 1,
432 resource: results_buffer.as_entire_binding(),
433 },
434 ],
435 });
436
437 let mut encoder = self
438 .device
439 .create_command_encoder(&wgpu::CommandEncoderDescriptor {
440 label: Some("Spacetime Compute Encoder"),
441 });
442
443 {
444 let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
445 label: Some("Spacetime Compute Pass"),
446 timestamp_writes: None,
447 });
448
449 compute_pass.set_pipeline(&self.spacetime_pipeline);
450 compute_pass.set_bind_group(0, &bind_group, &[]);
451
452 let workgroup_size = 64;
453 let num_workgroups = vectors.len().div_ceil(workgroup_size);
454 compute_pass.dispatch_workgroups(num_workgroups as u32, 1, 1);
455 }
456
457 let staging_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
459 label: Some("Staging Buffer"),
460 size: results_buffer.size(),
461 usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
462 mapped_at_creation: false,
463 });
464
465 encoder.copy_buffer_to_buffer(
466 &results_buffer,
467 0,
468 &staging_buffer,
469 0,
470 results_buffer.size(),
471 );
472
473 self.queue.submit([encoder.finish()]);
474
475 let buffer_slice = staging_buffer.slice(..);
476 let (sender, receiver) = futures::channel::oneshot::channel();
477 buffer_slice.map_async(wgpu::MapMode::Read, move |v| {
478 let _ = sender.send(v);
479 });
480
481 self.device.poll(wgpu::Maintain::wait()).panic_on_timeout();
482 receiver
483 .await
484 .map_err(|_| GpuError::BufferError("Failed to receive buffer mapping".to_string()))?
485 .map_err(|e| GpuError::BufferError(format!("Buffer mapping failed: {:?}", e)))?;
486
487 let data = buffer_slice.get_mapped_range();
488 let results: Vec<f32> = bytemuck::cast_slice(&data).to_vec();
489
490 drop(data);
491 staging_buffer.unmap();
492
493 Ok(results)
494 }
495
496 pub async fn propagate_particles(
498 &self,
499 particles: &[GpuRelativisticParticle],
500 params: &GpuTrajectoryParams,
501 ) -> Result<Vec<GpuRelativisticParticle>, GpuError> {
502 if particles.is_empty() || params.steps == 0 {
503 return Ok(particles.to_vec());
504 }
505 Self::validate_trajectory_params(params)?;
506 for (index, particle) in particles.iter().enumerate() {
507 Self::validate_particle(index, particle)?;
508 }
509
510 let particles_buffer = self
511 .device
512 .create_buffer_init(&wgpu::util::BufferInitDescriptor {
513 label: Some("Particles Buffer"),
514 contents: bytemuck::cast_slice(particles),
515 usage: wgpu::BufferUsages::STORAGE
516 | wgpu::BufferUsages::COPY_DST
517 | wgpu::BufferUsages::COPY_SRC,
518 });
519
520 let params_buffer = self
521 .device
522 .create_buffer_init(&wgpu::util::BufferInitDescriptor {
523 label: Some("Trajectory Params Buffer"),
524 contents: bytemuck::cast_slice(&[*params]),
525 usage: wgpu::BufferUsages::UNIFORM,
526 });
527
528 let bind_group_layout = self.geodesic_pipeline.get_bind_group_layout(0);
529 let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
530 label: Some("Geodesic Bind Group"),
531 layout: &bind_group_layout,
532 entries: &[
533 wgpu::BindGroupEntry {
534 binding: 0,
535 resource: particles_buffer.as_entire_binding(),
536 },
537 wgpu::BindGroupEntry {
538 binding: 1,
539 resource: params_buffer.as_entire_binding(),
540 },
541 ],
542 });
543
544 for _ in 0..params.steps {
546 let mut encoder = self
547 .device
548 .create_command_encoder(&wgpu::CommandEncoderDescriptor {
549 label: Some("Geodesic Compute Encoder"),
550 });
551
552 {
553 let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
554 label: Some("Geodesic Compute Pass"),
555 timestamp_writes: None,
556 });
557
558 compute_pass.set_pipeline(&self.geodesic_pipeline);
559 compute_pass.set_bind_group(0, &bind_group, &[]);
560
561 let workgroup_size = 64;
562 let num_workgroups = particles.len().div_ceil(workgroup_size);
563 compute_pass.dispatch_workgroups(num_workgroups as u32, 1, 1);
564 }
565
566 self.queue.submit([encoder.finish()]);
567 self.device.poll(wgpu::Maintain::wait()).panic_on_timeout();
568 }
569
570 let staging_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
572 label: Some("Particles Staging Buffer"),
573 size: particles_buffer.size(),
574 usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
575 mapped_at_creation: false,
576 });
577
578 let mut encoder = self
579 .device
580 .create_command_encoder(&wgpu::CommandEncoderDescriptor {
581 label: Some("Copy Encoder"),
582 });
583
584 encoder.copy_buffer_to_buffer(
585 &particles_buffer,
586 0,
587 &staging_buffer,
588 0,
589 particles_buffer.size(),
590 );
591 self.queue.submit([encoder.finish()]);
592
593 let buffer_slice = staging_buffer.slice(..);
594 let (sender, receiver) = futures::channel::oneshot::channel();
595 buffer_slice.map_async(wgpu::MapMode::Read, move |v| {
596 let _ = sender.send(v);
597 });
598
599 self.device.poll(wgpu::Maintain::wait()).panic_on_timeout();
600 receiver
601 .await
602 .map_err(|_| GpuError::BufferError("Failed to receive buffer mapping".to_string()))?
603 .map_err(|e| GpuError::BufferError(format!("Buffer mapping failed: {:?}", e)))?;
604
605 let data = buffer_slice.get_mapped_range();
606 let results: Vec<GpuRelativisticParticle> = bytemuck::cast_slice(&data).to_vec();
607
608 drop(data);
609 staging_buffer.unmap();
610
611 Ok(results)
612 }
613
614 fn validate_trajectory_params(params: &GpuTrajectoryParams) -> Result<(), GpuError> {
615 if !params.dt.is_finite() || params.dt <= 0.0 {
616 return Err(GpuError::BufferError(
617 "trajectory dt must be finite and positive".to_string(),
618 ));
619 }
620 if !params.tolerance.is_finite() || params.tolerance < 0.0 {
621 return Err(GpuError::BufferError(
622 "trajectory tolerance must be finite and non-negative".to_string(),
623 ));
624 }
625 if params.renorm_freq == 0 {
626 return Err(GpuError::BufferError(
627 "trajectory renorm_freq must be greater than zero".to_string(),
628 ));
629 }
630 if !params.schwarzschild_radius.is_finite() || params.schwarzschild_radius < 0.0 {
631 return Err(GpuError::BufferError(
632 "schwarzschild_radius must be finite and non-negative".to_string(),
633 ));
634 }
635 if !params.gm_parameter.is_finite() {
636 return Err(GpuError::BufferError(
637 "gm_parameter must be finite".to_string(),
638 ));
639 }
640
641 Ok(())
642 }
643
644 fn validate_particle(index: usize, particle: &GpuRelativisticParticle) -> Result<(), GpuError> {
645 if !particle.position.is_finite() {
646 return Err(GpuError::BufferError(format!(
647 "particle {index} position contains non-finite values"
648 )));
649 }
650 if !particle.velocity.is_finite() {
651 return Err(GpuError::BufferError(format!(
652 "particle {index} velocity contains non-finite values"
653 )));
654 }
655 if !particle.mass.is_finite() || particle.mass < 0.0 {
656 return Err(GpuError::BufferError(format!(
657 "particle {index} mass must be finite and non-negative"
658 )));
659 }
660 if !particle.charge.is_finite() {
661 return Err(GpuError::BufferError(format!(
662 "particle {index} charge must be finite"
663 )));
664 }
665 if !particle.proper_time.is_finite() {
666 return Err(GpuError::BufferError(format!(
667 "particle {index} proper_time must be finite"
668 )));
669 }
670
671 Ok(())
672 }
673}
674
675impl From<&RelativisticParticle> for GpuRelativisticParticle {
677 fn from(particle: &RelativisticParticle) -> Self {
678 let pos = &particle.position;
679 let vel = particle.four_velocity.as_spacetime_vector();
680
681 Self {
682 position: GpuSpacetimeVector::from_spacetime_vector(pos),
683 velocity: GpuSpacetimeVector::from_spacetime_vector(vel),
684 mass: particle.mass as f32,
685 charge: particle.charge as f32,
686 proper_time: 0.0, _padding: [0.0; 5],
688 }
689 }
690}
691
692#[cfg(test)]
693mod tests {
694 use super::*;
695
696 #[test]
697 fn test_gpu_spacetime_vector_conversion() {
698 let cpu_vector = SpacetimeVector::new(1.0, 2.0, 3.0, 4.0);
699 let gpu_vector = GpuSpacetimeVector::from_spacetime_vector(&cpu_vector);
700 let converted_back = gpu_vector.to_spacetime_vector();
701
702 assert!((converted_back.time() - 1.0).abs() < 1e-6);
703 assert_eq!(converted_back.x(), 2.0);
704 assert_eq!(converted_back.y(), 3.0);
705 assert_eq!(converted_back.z(), 4.0);
706 }
707
708 #[tokio::test]
709 #[ignore] async fn test_gpu_minkowski_products() {
711 let gpu_physics = match GpuRelativisticPhysics::new().await {
712 Ok(physics) => physics,
713 Err(_) => {
714 println!("GPU not available, skipping test");
715 return;
716 }
717 };
718
719 let vectors = vec![
720 GpuSpacetimeVector::new(1.0, 0.5, 0.0, 0.0),
721 GpuSpacetimeVector::new(2.0, 1.0, 0.0, 0.0),
722 ];
723
724 let results = gpu_physics
725 .compute_minkowski_products(&vectors)
726 .await
727 .unwrap();
728
729 assert_eq!(results.len(), vectors.len());
731
732 assert!((results[0] - (1.0 - 0.25)).abs() < 1e-6); assert!((results[1] - (4.0 - 1.0)).abs() < 1e-6); }
736}