Skip to main content

oxiphysics_python/world_api/
py_bindings.rs

1// Copyright 2026 COOLJAPAN OU (Team KitaSan)
2// SPDX-License-Identifier: Apache-2.0
3
4//! PyO3-style binding wrappers for MD, LBM, FEM, SPH simulations
5//! and material query helpers.
6
7#![allow(missing_docs)]
8
9use super::lbm::{PyLbmConfig, PyLbmGrid};
10use super::sph::{PySphConfig, PySphSim};
11
12// ===========================================================================
13// MD Simulation Bindings (PyO3-style API wrappers)
14// ===========================================================================
15
16use crate::md_api::{PyMdConfig, PyMdSimulation};
17
18/// PyO3-style binding handle for an MD simulation.
19///
20/// Wraps `PyMdSimulation` and exposes flat-array query helpers that mirror the
21/// `#[pymethods]` surface that would be used with the real PyO3 macro
22/// expansion.  All methods take and return primitive Rust types so that
23/// future `#[pyfunction]` / `#[pymethods]` annotations require no changes.
24#[derive(Debug, Clone)]
25#[allow(dead_code)]
26pub struct PyMdBinding {
27    sim: PyMdSimulation,
28}
29
30impl PyMdBinding {
31    // ------------------------------------------------------------------
32    // Construction
33    // ------------------------------------------------------------------
34
35    /// Create a new MD binding with the default (argon-reduced) configuration.
36    pub fn new_default() -> Self {
37        Self {
38            sim: PyMdSimulation::new(PyMdConfig::default()),
39        }
40    }
41
42    /// Create from explicit configuration.
43    pub fn new(config: PyMdConfig) -> Self {
44        Self {
45            sim: PyMdSimulation::new(config),
46        }
47    }
48
49    // ------------------------------------------------------------------
50    // Particle management
51    // ------------------------------------------------------------------
52
53    /// Add an atom at `[x, y, z]` with optional velocity `[vx, vy, vz]`.
54    ///
55    /// Returns the index of the newly added atom.
56    #[allow(clippy::too_many_arguments)]
57    pub fn add_atom(
58        &mut self,
59        x: f64,
60        y: f64,
61        z: f64,
62        vx: f64,
63        vy: f64,
64        vz: f64,
65        atom_type: u32,
66    ) -> usize {
67        let idx = self.sim.add_atom([x, y, z], atom_type);
68        self.sim.set_velocity(idx, [vx, vy, vz]);
69        idx
70    }
71
72    /// Return the number of atoms.
73    pub fn atom_count(&self) -> usize {
74        self.sim.atom_count()
75    }
76
77    // ------------------------------------------------------------------
78    // Stepping
79    // ------------------------------------------------------------------
80
81    /// Advance the simulation by `dt` (reduced units or SI, consistent with config).
82    pub fn step(&mut self, dt: f64) {
83        self.sim.step(dt);
84    }
85
86    /// Advance the simulation by `n_steps` steps, each of size `dt`.
87    pub fn run_steps(&mut self, n_steps: u32, dt: f64) {
88        for _ in 0..n_steps {
89            self.sim.step(dt);
90        }
91    }
92
93    // ------------------------------------------------------------------
94    // State queries
95    // ------------------------------------------------------------------
96
97    /// Return positions as a flat `[x0,y0,z0, x1,y1,z1, …]` array.
98    pub fn get_positions_flat(&self) -> Vec<f64> {
99        self.sim.all_positions()
100    }
101
102    /// Return velocities as a flat array.
103    pub fn get_velocities_flat(&self) -> Vec<f64> {
104        self.sim.all_velocities()
105    }
106
107    /// Return `[kinetic, potential, total]` energy.
108    pub fn get_energies(&self) -> [f64; 3] {
109        let ke = self.sim.kinetic_energy();
110        let pe = self.sim.potential_energy();
111        [ke, pe, ke + pe]
112    }
113
114    /// Instantaneous temperature in reduced units (from equipartition theorem).
115    pub fn get_temperature(&self) -> f64 {
116        self.sim.temperature()
117    }
118
119    /// Total simulation time elapsed.
120    pub fn elapsed_time(&self) -> f64 {
121        self.sim.time()
122    }
123
124    /// Number of completed time steps.
125    pub fn step_count(&self) -> u64 {
126        self.sim.step_count()
127    }
128
129    // ------------------------------------------------------------------
130    // Thermostat control
131    // ------------------------------------------------------------------
132
133    /// Enable or disable the velocity-rescaling thermostat at runtime.
134    pub fn set_thermostat_enabled(&mut self, enabled: bool) {
135        self.sim.set_thermostat(enabled);
136    }
137
138    /// Set the target temperature (reduced units).
139    pub fn set_target_temperature(&mut self, temp: f64) {
140        self.sim.set_target_temperature(temp);
141    }
142}
143
144// ===========================================================================
145// LBM Simulation Bindings (PyO3-style API wrappers)
146// ===========================================================================
147
148/// PyO3-style binding handle for a 2-D LBM simulation.
149///
150/// Uses the `PyLbmGrid` type (D2Q9 BGK) defined earlier in this file.
151#[derive(Debug, Clone)]
152#[allow(dead_code)]
153pub struct PyLbmBinding {
154    grid: PyLbmGrid,
155}
156
157impl PyLbmBinding {
158    // ------------------------------------------------------------------
159    // Construction
160    // ------------------------------------------------------------------
161
162    /// Create a new LBM binding with the given grid dimensions and viscosity.
163    pub fn new(width: usize, height: usize, viscosity: f64) -> Self {
164        let config = PyLbmConfig::new(width, height, viscosity);
165        Self {
166            grid: PyLbmGrid::new(&config),
167        }
168    }
169
170    /// Create from an explicit `PyLbmConfig`.
171    pub fn from_config(config: PyLbmConfig) -> Self {
172        Self {
173            grid: PyLbmGrid::new(&config),
174        }
175    }
176
177    // ------------------------------------------------------------------
178    // Grid dimensions
179    // ------------------------------------------------------------------
180
181    /// Grid width (number of cells along X).
182    pub fn grid_width(&self) -> usize {
183        self.grid.width()
184    }
185
186    /// Grid height (number of cells along Y).
187    pub fn grid_height(&self) -> usize {
188        self.grid.height()
189    }
190
191    /// Total number of cells.
192    pub fn cell_count(&self) -> usize {
193        self.grid.width() * self.grid.height()
194    }
195
196    // ------------------------------------------------------------------
197    // Stepping
198    // ------------------------------------------------------------------
199
200    /// Advance the simulation by one LBM step.
201    pub fn step(&mut self) {
202        self.grid.step();
203    }
204
205    /// Advance the simulation by `n_steps` steps.
206    pub fn run_steps(&mut self, n_steps: u32) {
207        for _ in 0..n_steps {
208            self.grid.step();
209        }
210    }
211
212    // ------------------------------------------------------------------
213    // Field queries
214    // ------------------------------------------------------------------
215
216    /// Return the density field as a flat row-major array.
217    pub fn get_density_flat(&self) -> Vec<f64> {
218        self.grid.density_field()
219    }
220
221    /// Return velocity field as interleaved `[ux0,uy0, ux1,uy1, …]`.
222    pub fn get_velocity_flat(&self) -> Vec<f64> {
223        self.grid.velocity_field()
224    }
225
226    /// Return velocity magnitude (speed) for each cell as a flat array.
227    pub fn get_speed_flat(&self) -> Vec<f64> {
228        let vel = self.grid.velocity_field();
229        vel.chunks(2)
230            .map(|c| (c[0] * c[0] + c[1] * c[1]).sqrt())
231            .collect()
232    }
233
234    /// Mean density across all cells.
235    pub fn mean_density(&self) -> f64 {
236        let rho = self.get_density_flat();
237        if rho.is_empty() {
238            return 0.0;
239        }
240        rho.iter().sum::<f64>() / rho.len() as f64
241    }
242
243    /// Maximum speed across all cells.
244    pub fn max_speed(&self) -> f64 {
245        self.get_speed_flat()
246            .iter()
247            .cloned()
248            .fold(f64::NEG_INFINITY, f64::max)
249            .max(0.0)
250    }
251
252    /// Number of completed steps.
253    pub fn step_count(&self) -> u64 {
254        self.grid.step_count()
255    }
256}
257
258// ===========================================================================
259// FEM Simulation Bindings (PyO3-style API wrappers)
260// ===========================================================================
261
262use crate::fem_api::{PyFemDirichletBC, PyFemMaterial, PyFemMesh, PyFemSolveResult, PyFemSolver};
263
264/// PyO3-style binding handle for a 2-D FEM solver.
265///
266/// Wraps `PyFemSolver` and `PyFemMesh` and exposes helpers suitable for
267/// future `#[pymethods]` annotation.
268#[derive(Debug, Clone)]
269#[allow(dead_code)]
270pub struct PyFemBinding {
271    mesh: PyFemMesh,
272    solver: PyFemSolver,
273    last_result: Option<PyFemSolveResult>,
274}
275
276impl PyFemBinding {
277    // ------------------------------------------------------------------
278    // Construction
279    // ------------------------------------------------------------------
280
281    /// Create a new FEM binding with an empty mesh.
282    pub fn new() -> Self {
283        Self {
284            mesh: PyFemMesh::new(),
285            solver: PyFemSolver::new(),
286            last_result: None,
287        }
288    }
289
290    // ------------------------------------------------------------------
291    // Mesh construction helpers
292    // ------------------------------------------------------------------
293
294    /// Add a 2-D node at position `(x, y)`.  Returns the new node index.
295    pub fn add_node(&mut self, x: f64, y: f64) -> usize {
296        self.mesh.add_node(x, y)
297    }
298
299    /// Add a CST triangular element given three node indices.
300    ///
301    /// Returns the new element index.
302    pub fn add_tri_element(&mut self, n0: usize, n1: usize, n2: usize) -> usize {
303        self.mesh.add_element(n0, n1, n2)
304    }
305
306    /// Add a CST triangular element with explicit material and thickness.
307    #[allow(clippy::too_many_arguments)]
308    pub fn add_tri_element_with_material(
309        &mut self,
310        n0: usize,
311        n1: usize,
312        n2: usize,
313        material_id: usize,
314        thickness: f64,
315    ) -> usize {
316        self.mesh
317            .add_element_with_material(n0, n1, n2, material_id, thickness)
318    }
319
320    /// Add a material to the mesh library. Returns the material index.
321    pub fn add_material(&mut self, material: PyFemMaterial) -> usize {
322        self.mesh.add_material(material)
323    }
324
325    /// Return the number of nodes.
326    pub fn node_count(&self) -> usize {
327        self.mesh.nodes.len()
328    }
329
330    /// Return the number of elements.
331    pub fn element_count(&self) -> usize {
332        self.mesh.elements.len()
333    }
334
335    // ------------------------------------------------------------------
336    // Boundary conditions and loads
337    // ------------------------------------------------------------------
338
339    /// Pin a node (fix both x and y DOFs to zero).
340    pub fn pin_node(&mut self, node: usize) {
341        self.mesh.pin_node(node);
342    }
343
344    /// Apply a Dirichlet (fixed DOF) boundary condition.
345    pub fn add_dirichlet_bc(&mut self, bc: PyFemDirichletBC) {
346        self.mesh.add_dirichlet_bc(bc);
347    }
348
349    /// Apply a nodal force `[fx, fy]` at `node`.
350    pub fn apply_nodal_force(&mut self, node: usize, fx: f64, fy: f64) {
351        self.mesh.apply_nodal_force(node, fx, fy);
352    }
353
354    // ------------------------------------------------------------------
355    // Solve
356    // ------------------------------------------------------------------
357
358    /// Solve the linear-elastic static problem.
359    ///
360    /// Returns `true` on success.  Query results via `get_displacement_flat`.
361    pub fn solve(&mut self) -> bool {
362        match self.solver.solve(&self.mesh) {
363            Some(result) => {
364                self.last_result = Some(result);
365                true
366            }
367            None => false,
368        }
369    }
370
371    /// Return nodal displacements as a flat `[ux0,uy0, ux1,uy1, …]`
372    /// array, or an empty `Vec` if no solve has been run.
373    pub fn get_displacement_flat(&self) -> Vec<f64> {
374        match &self.last_result {
375            Some(r) => r.displacements.clone(),
376            None => Vec::new(),
377        }
378    }
379
380    /// Return element von-Mises stresses as a flat array.
381    pub fn get_stress_flat(&self) -> Vec<f64> {
382        match &self.last_result {
383            Some(r) => r.von_mises_stress.clone(),
384            None => Vec::new(),
385        }
386    }
387
388    /// Return the maximum von-Mises stress, or `0.0` if no result exists.
389    pub fn max_von_mises_stress(&self) -> f64 {
390        match &self.last_result {
391            Some(r) if !r.von_mises_stress.is_empty() => r.max_von_mises().max(0.0),
392            _ => 0.0,
393        }
394    }
395
396    /// Return whether the last solve converged.
397    pub fn last_solve_converged(&self) -> bool {
398        self.last_result
399            .as_ref()
400            .map(|r| r.converged)
401            .unwrap_or(false)
402    }
403}
404
405impl Default for PyFemBinding {
406    fn default() -> Self {
407        Self::new()
408    }
409}
410
411// ===========================================================================
412// SPH Simulation Bindings (PyO3-style API wrappers)
413// ===========================================================================
414
415/// PyO3-style binding handle for a 3-D SPH simulation.
416///
417/// Uses the `PySphSim` type (WCSPH) defined earlier in this file.
418#[derive(Debug, Clone)]
419#[allow(dead_code)]
420pub struct PySphBinding {
421    sim: PySphSim,
422}
423
424impl PySphBinding {
425    // ------------------------------------------------------------------
426    // Construction
427    // ------------------------------------------------------------------
428
429    /// Create a new SPH binding with the water-like default configuration.
430    pub fn new_water() -> Self {
431        Self {
432            sim: PySphSim::new(PySphConfig::water()),
433        }
434    }
435
436    /// Create from an explicit `PySphConfig`.
437    pub fn from_config(config: PySphConfig) -> Self {
438        Self {
439            sim: PySphSim::new(config),
440        }
441    }
442
443    // ------------------------------------------------------------------
444    // Particle management
445    // ------------------------------------------------------------------
446
447    /// Add a fluid particle at `[x, y, z]` with zero initial velocity.
448    ///
449    /// Returns the new particle index.
450    pub fn add_particle_at(&mut self, x: f64, y: f64, z: f64) -> usize {
451        let idx = self.sim.particle_count();
452        self.sim.add_particle([x, y, z]);
453        idx
454    }
455
456    /// Add a fluid particle at `[x, y, z]` with the given initial velocity
457    /// (`vx`, `vy`, `vz`).
458    pub fn add_particle(&mut self, x: f64, y: f64, z: f64, vx: f64, vy: f64, vz: f64) -> usize {
459        let idx = self.sim.particle_count();
460        self.sim.add_particle([x, y, z]);
461        // Set velocity via the position-then-velocity pattern
462        if let Some(vel) = self.sim.velocities_mut().get_mut(idx) {
463            *vel = [vx, vy, vz];
464        }
465        idx
466    }
467
468    /// Return the number of particles.
469    pub fn particle_count(&self) -> usize {
470        self.sim.particle_count()
471    }
472
473    // ------------------------------------------------------------------
474    // Stepping
475    // ------------------------------------------------------------------
476
477    /// Advance the simulation by `dt` seconds.
478    pub fn step(&mut self, dt: f64) {
479        self.sim.step(dt);
480    }
481
482    /// Advance by `n_steps` steps.
483    pub fn run_steps(&mut self, n_steps: u32, dt: f64) {
484        for _ in 0..n_steps {
485            self.sim.step(dt);
486        }
487    }
488
489    // ------------------------------------------------------------------
490    // State queries
491    // ------------------------------------------------------------------
492
493    /// Return positions as a flat `[x0,y0,z0, x1,y1,z1, …]` array.
494    pub fn get_positions_flat(&self) -> Vec<f64> {
495        self.sim.all_positions()
496    }
497
498    /// Return velocities as a flat array.
499    pub fn get_velocities_flat(&self) -> Vec<f64> {
500        self.sim.all_velocities()
501    }
502
503    /// Return densities as a flat array (one value per particle).
504    pub fn get_densities_flat(&self) -> Vec<f64> {
505        (0..self.sim.particle_count())
506            .filter_map(|i| self.sim.density(i))
507            .collect()
508    }
509
510    /// Return pressures as a flat array.
511    pub fn get_pressures_flat(&self) -> Vec<f64> {
512        (0..self.sim.particle_count())
513            .filter_map(|i| self.sim.pressure(i))
514            .collect()
515    }
516
517    /// Return speeds (velocity magnitudes) as a flat array.
518    pub fn get_speeds_flat(&self) -> Vec<f64> {
519        let vels = self.sim.all_velocities();
520        vels.chunks(3)
521            .map(|c| (c[0] * c[0] + c[1] * c[1] + c[2] * c[2]).sqrt())
522            .collect()
523    }
524
525    /// Mean density of all particles.
526    pub fn mean_density(&self) -> f64 {
527        let n = self.sim.particle_count();
528        if n == 0 {
529            return 0.0;
530        }
531        let sum: f64 = (0..n).filter_map(|i| self.sim.density(i)).sum();
532        sum / n as f64
533    }
534
535    /// Maximum speed across all particles.
536    pub fn max_speed(&self) -> f64 {
537        self.get_speeds_flat()
538            .iter()
539            .cloned()
540            .fold(f64::NEG_INFINITY, f64::max)
541            .max(0.0)
542    }
543
544    /// Total kinetic energy of all particles.
545    pub fn total_kinetic_energy(&self) -> f64 {
546        let mass = self.sim.particle_mass();
547        let vels = self.sim.all_velocities();
548        vels.chunks(3)
549            .map(|c| 0.5 * mass * (c[0] * c[0] + c[1] * c[1] + c[2] * c[2]))
550            .sum()
551    }
552
553    // ------------------------------------------------------------------
554    // Configuration queries
555    // ------------------------------------------------------------------
556
557    /// Smoothing length h.
558    pub fn smoothing_length(&self) -> f64 {
559        self.sim.smoothing_length()
560    }
561
562    /// Rest density ρ₀.
563    pub fn rest_density(&self) -> f64 {
564        self.sim.rest_density()
565    }
566}
567
568// ===========================================================================
569// Material query helpers (PyO3-style standalone functions)
570// ===========================================================================
571
572/// Material property record returned by query functions.
573#[derive(Debug, Clone, PartialEq)]
574#[allow(dead_code)]
575pub struct MaterialProperties {
576    /// Material name.
577    pub name: String,
578    /// Density in kg/m³.
579    pub density: f64,
580    /// Young's modulus in Pa.
581    pub youngs_modulus: f64,
582    /// Poisson's ratio (dimensionless).
583    pub poisson_ratio: f64,
584    /// Ultimate tensile strength in Pa.
585    pub tensile_strength: f64,
586    /// Dynamic viscosity in Pa·s (0 for solids).
587    pub viscosity: f64,
588    /// Yield strength in Pa (0 if not applicable).
589    pub yield_strength: f64,
590}
591
592impl MaterialProperties {
593    /// Steel (structural grade).
594    pub fn steel() -> Self {
595        Self {
596            name: "steel".into(),
597            density: 7850.0,
598            youngs_modulus: 200e9,
599            poisson_ratio: 0.3,
600            tensile_strength: 400e6,
601            viscosity: 0.0,
602            yield_strength: 250e6,
603        }
604    }
605
606    /// Aluminium alloy 6061-T6.
607    pub fn aluminium() -> Self {
608        Self {
609            name: "aluminium".into(),
610            density: 2700.0,
611            youngs_modulus: 69e9,
612            poisson_ratio: 0.33,
613            tensile_strength: 310e6,
614            viscosity: 0.0,
615            yield_strength: 276e6,
616        }
617    }
618
619    /// Concrete (typical structural).
620    pub fn concrete() -> Self {
621        Self {
622            name: "concrete".into(),
623            density: 2400.0,
624            youngs_modulus: 30e9,
625            poisson_ratio: 0.2,
626            tensile_strength: 3e6,
627            viscosity: 0.0,
628            yield_strength: 0.0,
629        }
630    }
631
632    /// Natural rubber.
633    pub fn rubber() -> Self {
634        Self {
635            name: "rubber".into(),
636            density: 1100.0,
637            youngs_modulus: 0.01e9,
638            poisson_ratio: 0.49,
639            tensile_strength: 20e6,
640            viscosity: 0.0,
641            yield_strength: 0.0,
642        }
643    }
644
645    /// Water at 20 °C.
646    pub fn water() -> Self {
647        Self {
648            name: "water".into(),
649            density: 998.0,
650            youngs_modulus: 0.0,
651            poisson_ratio: 0.0,
652            tensile_strength: 0.0,
653            viscosity: 1.002e-3,
654            yield_strength: 0.0,
655        }
656    }
657
658    /// Air at 20 °C, 1 atm.
659    pub fn air() -> Self {
660        Self {
661            name: "air".into(),
662            density: 1.204,
663            youngs_modulus: 0.0,
664            poisson_ratio: 0.0,
665            tensile_strength: 0.0,
666            viscosity: 1.825e-5,
667            yield_strength: 0.0,
668        }
669    }
670
671    /// Titanium Ti-6Al-4V.
672    pub fn titanium() -> Self {
673        Self {
674            name: "titanium".into(),
675            density: 4430.0,
676            youngs_modulus: 114e9,
677            poisson_ratio: 0.34,
678            tensile_strength: 950e6,
679            viscosity: 0.0,
680            yield_strength: 880e6,
681        }
682    }
683
684    /// Lookup a material by name (case-insensitive).
685    ///
686    /// Returns `None` if the name is unknown.
687    pub fn lookup(name: &str) -> Option<Self> {
688        match name.to_lowercase().as_str() {
689            "steel" => Some(Self::steel()),
690            "aluminium" | "aluminum" => Some(Self::aluminium()),
691            "concrete" => Some(Self::concrete()),
692            "rubber" => Some(Self::rubber()),
693            "water" => Some(Self::water()),
694            "air" => Some(Self::air()),
695            "titanium" => Some(Self::titanium()),
696            _ => None,
697        }
698    }
699
700    /// Compute the speed of longitudinal acoustic wave (P-wave) in m/s.
701    ///
702    /// Returns `None` for fluids where `youngs_modulus == 0`.
703    pub fn p_wave_speed(&self) -> Option<f64> {
704        if self.youngs_modulus < 1.0 || self.density < 1e-6 {
705            return None;
706        }
707        let nu = self.poisson_ratio;
708        let e = self.youngs_modulus;
709        let rho = self.density;
710        let c = ((e * (1.0 - nu)) / (rho * (1.0 + nu) * (1.0 - 2.0 * nu))).sqrt();
711        Some(c)
712    }
713
714    /// Compute the shear wave speed (S-wave) in m/s.
715    ///
716    /// Returns `None` for fluids.
717    pub fn s_wave_speed(&self) -> Option<f64> {
718        if self.youngs_modulus < 1.0 || self.density < 1e-6 {
719            return None;
720        }
721        let nu = self.poisson_ratio;
722        let e = self.youngs_modulus;
723        let rho = self.density;
724        let g = e / (2.0 * (1.0 + nu));
725        Some((g / rho).sqrt())
726    }
727}
728
729// PyO3 standalone function mirrors (named with `py_` prefix to avoid future
730// name conflicts when the real `#[pyfunction]` wrappers are generated).
731
732/// PyO3-callable: look up a material and return `[density, E, nu, UTS, visc, yield]`
733/// or an empty vec if the name is unknown.
734#[allow(dead_code)]
735pub fn py_query_material(name: &str) -> Vec<f64> {
736    match MaterialProperties::lookup(name) {
737        Some(m) => vec![
738            m.density,
739            m.youngs_modulus,
740            m.poisson_ratio,
741            m.tensile_strength,
742            m.viscosity,
743            m.yield_strength,
744        ],
745        None => Vec::new(),
746    }
747}
748
749/// PyO3-callable: return P-wave speed for a named material, or `0.0`.
750#[allow(dead_code)]
751pub fn py_p_wave_speed(name: &str) -> f64 {
752    MaterialProperties::lookup(name)
753        .and_then(|m| m.p_wave_speed())
754        .unwrap_or(0.0)
755}
756
757/// PyO3-callable: return S-wave speed for a named material, or `0.0`.
758#[allow(dead_code)]
759pub fn py_s_wave_speed(name: &str) -> f64 {
760    MaterialProperties::lookup(name)
761        .and_then(|m| m.s_wave_speed())
762        .unwrap_or(0.0)
763}