vox_geometry_rust 0.1.2

/*
 * // Copyright (c) 2021 Feng Yang
 * //
 * // I am making my contributions/submissions to this project solely in my
 * // personal capacity and am not conveying any rights to any intellectual
 * // property of any third parties.
 */

use crate::vector3::Vector3D;
use crate::particle_system_solver3::ParticleSystemSolver3;
use crate::sph_system_data3::*;
use crate::collider3::Collider3Ptr;
use crate::particle_emitter3::ParticleEmitter3Ptr;
use crate::vector_field3::VectorField3Ptr;
use crate::sph_kernels3::SphSpikyKernel3;
use crate::animation::*;
use crate::physics_animation::*;
use std::sync::{RwLock, Arc};
use std::time::SystemTime;
use log::info;
use rayon::prelude::*;

///
/// # 3-D SPH solver.
///
/// This class implements 3-D SPH solver. The main pressure solver is based on
/// equation-of-state (EOS).
///
/// \see M{\"u}ller et al., Particle-based fluid simulation for interactive
///      applications, SCA 3003.
/// \see M. Becker and M. Teschner, Weakly compressible SPH for free surface
///      flows, SCA 3007.
/// \see Adams and Wicke, Meshless approximation methods and applications in
///      physics based modeling and animation, Eurographics tutorials 3009.
///
pub struct SphSolver3 {
    /// Exponent component of equation-of-state (or Tait's equation).
    _eos_exponent: f64,

    /// Negative pressure scaling factor.
    /// Zero means clamping. One means do nothing.
    _negative_pressure_scale: f64,

    /// Viscosity coefficient.
    _viscosity_coefficient: f64,

    /// Pseudo-viscosity coefficient velocity filtering.
    /// This is a minimum "safety-net" for SPH solver which is quite
    /// sensitive to the parameters.
    _pseudo_viscosity_coefficient: f64,

    /// Speed of sound in medium to determine the stiffness of the system.
    /// Ideally, it should be the actual speed of sound in the fluid, but in
    /// practice, use lower value to trace-off performance and compressibility.
    _speed_of_sound: f64,

    /// Scales the max allowed time-step.
    _time_step_limit_scale: f64,

    _base_solver: ParticleSystemSolver3,
    _particle_system_data: SphSystemData3Ptr,
}

impl SphSolver3 {
    /// Constructs a solver with empty particle set.
    pub fn new_default() -> SphSolver3 {
        let mut solver = SphSolver3 {
            _eos_exponent: 7.0,
            _negative_pressure_scale: 0.0,
            _viscosity_coefficient: 0.01,
            _pseudo_viscosity_coefficient: 10.0,
            _speed_of_sound: 100.0,
            _time_step_limit_scale: 1.0,
            _base_solver: ParticleSystemSolver3::new_default(),
            _particle_system_data: SphSystemData3Ptr::new(RwLock::new(SphSystemData3::new_default())),
        };
        solver.set_is_using_fixed_sub_time_steps(false);
        return solver;
    }

    /// Constructs a solver with target density, spacing, and relative kernel radius.
    pub fn new(target_density: f64,
               target_spacing: f64,
               relative_kernel_radius: f64) -> SphSolver3 {
        let sph_particles = SphSystemData3Ptr::new(RwLock::new(SphSystemData3::new_default()));
        sph_particles.write().unwrap().set_target_density(target_density);
        sph_particles.write().unwrap().set_target_spacing(target_spacing);
        sph_particles.write().unwrap().set_relative_kernel_radius(relative_kernel_radius);

        let mut solver = SphSolver3 {
            _eos_exponent: 7.0,
            _negative_pressure_scale: 0.0,
            _viscosity_coefficient: 0.01,
            _pseudo_viscosity_coefficient: 10.0,
            _speed_of_sound: 100.0,
            _time_step_limit_scale: 1.0,
            _base_solver: ParticleSystemSolver3::new_default(),
            _particle_system_data: sph_particles,
        };
        solver.set_is_using_fixed_sub_time_steps(false);
        return solver;
    }
}

impl SphSolver3 {
    /// Returns the drag coefficient.
    pub fn drag_coefficient(&self) -> f64 {
        return self._base_solver.drag_coefficient();
    }

    ///
    /// \brief      Sets the drag coefficient.
    ///
    /// The drag coefficient controls the amount of air-drag. The coefficient
    /// should be a positive number and 0 means no drag force.
    ///
    /// - parameter:  new_drag_coefficient The new drag coefficient.
    ///
    pub fn set_drag_coefficient(&mut self, new_drag_coefficient: f64) {
        self._base_solver.set_drag_coefficient(f64::max(new_drag_coefficient, 0.0));
    }

    /// Sets the restitution coefficient.
    pub fn restitution_coefficient(&self) -> f64 {
        return self._base_solver.restitution_coefficient();
    }

    ///
    /// \brief      Sets the restitution coefficient.
    ///
    /// The restitution coefficient controls the bouncy-ness of a particle when
    /// it hits a collider surface. The range of the coefficient should be 0 to
    /// 1 -- 0 means no bounce back and 1 means perfect reflection.
    ///
    /// - parameter:  new_restitution_coefficient The new restitution coefficient.
    ///
    pub fn set_restitution_coefficient(&mut self, new_restitution_coefficient: f64) {
        self._base_solver.set_restitution_coefficient(crate::math_utils::clamp(new_restitution_coefficient, 0.0, 1.0));
    }

    /// Returns the gravity.
    pub fn gravity(&self) -> &Vector3D {
        return self._base_solver.gravity();
    }

    /// Sets the gravity.
    pub fn set_gravity(&mut self, new_gravity: &Vector3D) {
        self._base_solver.set_gravity(new_gravity);
    }

    ///
    /// \brief      Returns the particle system data.
    ///
    /// This function returns the particle system data. The data is created when
    /// this solver is constructed and also owned by the solver.
    ///
    /// \return     The particle system data.
    ///
    pub fn particle_system_data(&self) -> &SphSystemData3Ptr {
        return &self._particle_system_data;
    }

    /// Returns the collider.
    pub fn collider(&self) -> &Option<Collider3Ptr> {
        return self._base_solver.collider();
    }

    /// Sets the collider.
    pub fn set_collider(&mut self, new_collider: &Collider3Ptr) {
        self._base_solver.set_collider(new_collider);
    }

    /// Returns the emitter.
    pub fn emitter(&self) -> &Option<ParticleEmitter3Ptr> {
        return self._base_solver.emitter();
    }

    /// Sets the emitter.
    pub fn set_emitter(&mut self, new_emitter: &ParticleEmitter3Ptr) {
        self._base_solver.set_emitter(new_emitter);
    }

    /// Returns the wind field.
    pub fn wind(&self) -> &VectorField3Ptr {
        return self._base_solver.wind();
    }

    ///
    /// \brief      Sets the wind.
    ///
    /// Wind can be applied to the particle system by setting a vector field to
    /// the solver.
    ///
    /// - parameter:  new_wind The new wind.
    ///
    pub fn set_wind(&mut self, new_wind: &VectorField3Ptr) {
        self._base_solver.set_wind(new_wind);
    }
}

impl SphSolver3 {
    /// Returns the exponent part of the equation-of-state.
    pub fn eos_exponent(&self) -> f64 {
        return self._eos_exponent;
    }
    ///
    /// \brief Sets the exponent part of the equation-of-state.
    ///
    /// This function sets the exponent part of the equation-of-state.
    /// The value must be greater than 1.0, and smaller inputs will be clamped.
    /// Default is 7.
    ///
    pub fn set_eos_exponent(&mut self, new_eos_exponent: f64) {
        self._eos_exponent = f64::max(new_eos_exponent, 1.0);
    }

    /// Returns the negative pressure scale.
    pub fn negative_pressure_scale(&self) -> f64 {
        return self._negative_pressure_scale;
    }

    ///
    /// \brief Sets the negative pressure scale.
    ///
    /// This function sets the negative pressure scale. By setting the number
    /// between 0 and 1, the solver will scale the effect of negative pressure
    /// which can prevent the clumping of the particles near the surface. Input
    /// value outside 0 and 1 will be clamped within the range. Default is 0.
    ///
    pub fn set_negative_pressure_scale(&mut self, new_negative_pressure_scale: f64) {
        self._negative_pressure_scale = crate::math_utils::clamp(new_negative_pressure_scale, 0.0, 1.0);
    }

    /// Returns the viscosity coefficient.
    pub fn viscosity_coefficient(&self) -> f64 {
        return self._viscosity_coefficient;
    }

    /// Sets the viscosity coefficient.
    pub fn set_viscosity_coefficient(&mut self, new_viscosity_coefficient: f64) {
        self._viscosity_coefficient = f64::max(new_viscosity_coefficient, 0.0);
    }

    /// Returns the pseudo viscosity coefficient.
    pub fn pseudo_viscosity_coefficient(&self) -> f64 {
        return self._pseudo_viscosity_coefficient;
    }

    ///
    /// \brief Sets the pseudo viscosity coefficient.
    ///
    /// This function sets the pseudo viscosity coefficient which applies
    /// additional pseudo-physical damping to the system. Default is 10.
    ///
    pub fn set_pseudo_viscosity_coefficient(&mut self, new_pseudo_viscosity_coefficient: f64) {
        self._pseudo_viscosity_coefficient = f64::max(new_pseudo_viscosity_coefficient, 0.0);
    }

    /// Returns the speed of sound.
    pub fn speed_of_sound(&self) -> f64 {
        return self._speed_of_sound;
    }

    ///
    /// \brief Sets the speed of sound.
    ///
    /// This function sets the speed of sound which affects the stiffness of the
    /// EOS and the time-step size. Higher value will make EOS stiffer and the
    /// time-step smaller. The input value must be higher than 0.0.
    ///
    pub fn set_speed_of_sound(&mut self, new_speed_of_sound: f64) {
        self._speed_of_sound = f64::max(new_speed_of_sound, f64::EPSILON);
    }

    ///
    /// \brief Multiplier that scales the max allowed time-step.
    ///
    /// This function returns the multiplier that scales the max allowed
    /// time-step. When the scale is 1.0, the time-step is bounded by the speed
    /// of sound and max acceleration.
    ///
    pub fn time_step_limit_scale(&self) -> f64 {
        return self._time_step_limit_scale;
    }

    ///
    /// \brief Sets the multiplier that scales the max allowed time-step.
    ///
    /// This function sets the multiplier that scales the max allowed
    /// time-step. When the scale is 1.0, the time-step is bounded by the speed
    /// of sound and max acceleration.
    ///
    pub fn set_time_step_limit_scale(&mut self, new_scale: f64) {
        self._time_step_limit_scale = f64::max(new_scale, 0.0);
    }
}

//--------------------------------------------------------------------------------------------------
impl Animation for SphSolver3 {
    fn on_update(&mut self, frame: &Frame) {
        PhysicsAnimation::on_update(self, frame);
    }
}

impl PhysicsAnimation for SphSolver3 {
    fn on_advance_time_step(&mut self, time_step_in_seconds: f64) {
        self.begin_advance_time_step(time_step_in_seconds);

        let mut timer = SystemTime::now();
        self.accumulate_forces(time_step_in_seconds);
        info!("Accumulating forces took {} seconds", timer.elapsed().unwrap().as_secs_f64());

        timer = SystemTime::now();
        self.time_integration(time_step_in_seconds);
        info!("Time integration took {} seconds", timer.elapsed().unwrap().as_secs_f64());

        timer = SystemTime::now();
        self.resolve_collision();
        info!("Resolving collision took {} seconds", timer.elapsed().unwrap().as_secs_f64());

        self.end_advance_time_step(time_step_in_seconds);
    }

    /// Returns the number of sub-time-steps.
    fn number_of_sub_time_steps(&self,
                                time_interval_in_seconds: f64) -> usize {
        let number_of_particles = self._particle_system_data.read().unwrap().number_of_particles();

        let kernel_radius = self._particle_system_data.read().unwrap().kernel_radius();
        let mass = self._particle_system_data.read().unwrap().mass();

        let mut max_force_magnitude = 0.0;

        for i in 0..number_of_particles {
            max_force_magnitude = f64::max(max_force_magnitude,
                                           self._particle_system_data.read().unwrap().forces()[i].length());
        }

        let time_step_limit_by_speed = K_TIME_STEP_LIMIT_BY_SPEED_FACTOR * kernel_radius / self._speed_of_sound;
        let time_step_limit_by_force = K_TIME_STEP_LIMIT_BY_FORCE_FACTOR * f64::sqrt(kernel_radius * mass / max_force_magnitude);

        let desired_time_step = self._time_step_limit_scale * f64::min(time_step_limit_by_speed, time_step_limit_by_force);

        return f64::ceil(time_interval_in_seconds / desired_time_step) as usize;
    }

    fn on_initialize(&mut self) {
        // When initializing the solver, update the collider and emitter state as
        // well since they also affects the initial condition of the simulation.
        let mut timer = SystemTime::now();
        self.update_collider(0.0);
        info!("Update collider took {} seconds", timer.elapsed().unwrap().as_secs_f64());

        timer = SystemTime::now();
        self.update_emitter(0.0);
        info!("Update emitter took {} seconds", timer.elapsed().unwrap().as_secs_f64());
    }

    fn view(&self) -> &PhysicsAnimationData {
        return self._base_solver.view();
    }

    fn view_mut(&mut self) -> &mut PhysicsAnimationData {
        return self._base_solver.view_mut();
    }
}

impl SphSolver3 {
    /// Resolves any collisions occurred by the particles.
    fn resolve_collision(&mut self) {
        if let Some(collider) = &self._base_solver.collider() {
            let collider = collider.clone();
            let radius = self._particle_system_data.read().unwrap().radius();
            let restitution = self._base_solver.restitution_coefficient();
            (&mut self._base_solver._new_positions, &mut self._base_solver._new_velocities).into_par_iter().for_each(|(pos, vel)| {
                collider.read().unwrap().resolve_collision(
                    radius,
                    restitution,
                    pos,
                    vel);
            });
        }
    }

    fn begin_advance_time_step(&mut self, time_step_in_seconds: f64) {
        // Clear forces
        self._particle_system_data.write().unwrap().forces_mut().fill(Vector3D::new_default());

        // Update collider and emitter
        let mut timer = SystemTime::now();
        self.update_collider(time_step_in_seconds);
        info!("Update collider took {} seconds", timer.elapsed().unwrap().as_secs_f64());

        timer = SystemTime::now();
        self.update_emitter(time_step_in_seconds);
        info!("Update emitter took {} seconds", timer.elapsed().unwrap().as_secs_f64());

        // Allocate buffers
        let n = self._particle_system_data.read().unwrap().number_of_particles();
        self._base_solver._new_positions.resize(n, Vector3D::new_default());
        self._base_solver._new_velocities.resize(n, Vector3D::new_default());

        self.on_begin_advance_time_step(time_step_in_seconds);
    }

    fn end_advance_time_step(&mut self, time_step_in_seconds: f64) {
        // Update data
        (self._particle_system_data.write().unwrap().positions_mut(),
         self._particle_system_data.write().unwrap().velocities_mut(),
         &self._base_solver._new_positions, &self._base_solver._new_velocities).into_par_iter().for_each(|(pos, vel, pos_new, vel_new)| {
            *pos = *pos_new;
            *vel = *vel_new;
        });

        self.on_end_advance_time_step(time_step_in_seconds);
    }

    fn accumulate_external_forces(&mut self) {
        let mass = self._particle_system_data.read().unwrap().mass();

        (self._particle_system_data.write().unwrap().forces_mut(),
         self._particle_system_data.read().unwrap().positions(),
         self._particle_system_data.read().unwrap().velocities()).into_par_iter().for_each(|(force, pos, vel)| {
            // Gravity
            let mut f = self._base_solver._gravity * mass;

            // Wind forces
            let relative_vel = *vel - self._base_solver._wind.read().unwrap().sample(pos);
            f += relative_vel * -self._base_solver._drag_coefficient;

            *force += f;
        });
    }

    fn time_integration(&mut self, time_step_in_seconds: f64) {
        let mass = self._particle_system_data.read().unwrap().mass();

        (&mut self._base_solver._new_positions, &mut self._base_solver._new_velocities,
         self._particle_system_data.read().unwrap().forces(),
         self._particle_system_data.read().unwrap().positions(),
         self._particle_system_data.read().unwrap().velocities()).into_par_iter().for_each(|(new_pos, new_vel, force, pos, vel)| {
            // Integrate velocity first
            *new_vel = *vel + *force * time_step_in_seconds / mass;

            // Integrate position.
            *new_pos = *pos + *new_vel * time_step_in_seconds;
        });
    }

    fn update_collider(&mut self, time_step_in_seconds: f64) {
        if let Some(collider) = &self._base_solver._collider {
            collider.write().unwrap().update(self.current_time_in_seconds(),
                                             time_step_in_seconds);
        }
    }

    fn update_emitter(&mut self, time_step_in_seconds: f64) {
        if let Some(emitter) = &self._base_solver._emitter {
            emitter.write().unwrap().update(self.current_time_in_seconds(),
                                            time_step_in_seconds);
        }
    }
}

const K_TIME_STEP_LIMIT_BY_SPEED_FACTOR: f64 = 0.4;
const K_TIME_STEP_LIMIT_BY_FORCE_FACTOR: f64 = 0.35;

impl SphSolver3 {
    /// Accumulates the force to the forces array in the particle system.
    fn accumulate_forces(&mut self, time_step_in_seconds: f64) {
        self.accumulate_non_pressure_forces(time_step_in_seconds);
        self.accumulate_pressure_force(time_step_in_seconds);
    }

    /// Performs pre-processing step before the simulation.
    fn on_begin_advance_time_step(&self, _: f64) {
        let timer = SystemTime::now();
        self._particle_system_data.write().unwrap().build_neighbor_searcher();
        self._particle_system_data.write().unwrap().build_neighbor_lists();
        self._particle_system_data.write().unwrap().update_densities();

        info!("Building neighbor lists and updating densities took {} seconds", timer.elapsed().unwrap().as_secs_f64());
    }

    /// Performs post-processing step before the simulation.
    fn on_end_advance_time_step(&self, time_step_in_seconds: f64) {
        self.compute_pseudo_viscosity(time_step_in_seconds);

        let number_of_particles = self._particle_system_data.read().unwrap().number_of_particles();

        let mut max_density = 0.0;
        for i in 0..number_of_particles {
            max_density = f64::max(max_density, self._particle_system_data.read().unwrap().densities()[i]);
        }

        info!("Max density: {} Max density / target density ratio: {}", max_density,
              max_density / self._particle_system_data.read().unwrap().target_density());
    }

    /// Accumulates the non-pressure forces to the forces array in the particle
    /// system.
    fn accumulate_non_pressure_forces(&mut self, _: f64) {
        // Add external forces
        self.accumulate_external_forces();
        self.accumulate_viscosity_force();
    }

    /// Accumulates the pressure force to the forces array in the particle
    /// system.
    fn accumulate_pressure_force(&self, _: f64) {
        self.compute_pressure();
        self.accumulate_pressure_force_new(
            self._particle_system_data.read().unwrap().positions(),
            self._particle_system_data.read().unwrap().densities(),
            self._particle_system_data.read().unwrap().pressures(),
            self._particle_system_data.write().unwrap().forces_mut());
    }

    /// Computes the pressure.
    fn compute_pressure(&self) {
        // See Murnaghan-Tait equation of state from
        // https://en.wikipedia.org/wiki/Tait_equation
        let target_density = self._particle_system_data.read().unwrap().target_density();
        let eos_scale = target_density * crate::math_utils::square(self._speed_of_sound);

        (self._particle_system_data.write().unwrap().pressures_mut(),
         self._particle_system_data.read().unwrap().densities()).into_par_iter().for_each(|(p, d)| {
            *p = crate::physics_helpers::compute_pressure_from_eos(
                *d,
                target_density,
                eos_scale,
                self.eos_exponent(),
                self.negative_pressure_scale());
        });
    }

    /// Accumulates the pressure force to the given \p pressure_forces array.
    fn accumulate_pressure_force_new(&self, positions: &Vec<Vector3D>,
                                     densities: &Vec<f64>,
                                     pressures: &Vec<f64>,
                                     pressure_forces: &mut Vec<Vector3D>) {
        let number_of_particles = self._particle_system_data.read().unwrap().number_of_particles();

        let mass_squared = crate::math_utils::square(self._particle_system_data.read().unwrap().mass());
        let kernel = SphSpikyKernel3::new(self._particle_system_data.read().unwrap().kernel_radius());

        (pressure_forces,
         self._particle_system_data.read().unwrap().neighbor_lists(),
         0..number_of_particles).into_par_iter().for_each(|(f, neighbors, index)| {
            for j in neighbors {
                let dist = positions[index].distance_to(positions[*j]);

                if dist > 0.0 {
                    let dir = (positions[*j] - positions[index]) / dist;
                    *f -= kernel.gradient_dir(dist, &dir)
                        * mass_squared
                        * (pressures[index] / (densities[index] * densities[index]) + pressures[*j] / (densities[*j] * densities[*j]));
                }
            }
        });
    }

    /// Accumulates the viscosity force to the forces array in the particle
    /// system.
    fn accumulate_viscosity_force(&self) {
        let number_of_particles = self._particle_system_data.read().unwrap().number_of_particles();
        let x = self._particle_system_data.read().unwrap().positions().clone();
        let v = self._particle_system_data.read().unwrap().velocities().clone();
        let d = self._particle_system_data.read().unwrap().densities().clone();

        let mass_squared = crate::math_utils::square(self._particle_system_data.read().unwrap().mass());
        let kernel = SphSpikyKernel3::new(self._particle_system_data.read().unwrap().kernel_radius());

        (self._particle_system_data.write().unwrap().forces_mut(),
         self._particle_system_data.read().unwrap().neighbor_lists(),
         0..number_of_particles).into_par_iter().for_each(|(f, neighbors, index)| {
            for j in neighbors {
                let dist = x[index].distance_to(x[*j]);

                *f += (v[*j] - v[index]) * self.viscosity_coefficient() * mass_squared / d[*j]
                    * kernel.second_derivative(dist);
            }
        });
    }

    /// Computes pseudo viscosity.
    fn compute_pseudo_viscosity(&self, time_step_in_seconds: f64) {
        let number_of_particles = self._particle_system_data.read().unwrap().number_of_particles();
        let x = self._particle_system_data.read().unwrap().positions().clone();
        let v = self._particle_system_data.read().unwrap().velocities().clone();
        let d = self._particle_system_data.read().unwrap().densities().clone();

        let mass = self._particle_system_data.read().unwrap().mass();
        let kernel = SphSpikyKernel3::new(self._particle_system_data.read().unwrap().kernel_radius());

        let mut smoothed_velocities: Vec<Vector3D> = Vec::new();
        smoothed_velocities.resize(number_of_particles, Vector3D::new_default());

        (&mut smoothed_velocities,
         self._particle_system_data.read().unwrap().neighbor_lists(),
         0..number_of_particles).into_par_iter().for_each(|(vel, neighbors, index)| {
            let mut weight_sum = 0.0;
            let mut smoothed_velocity = Vector3D::new_default();

            for j in neighbors {
                let dist = x[index].distance_to(x[*j]);
                let wj = mass / d[*j] * kernel.apply(dist);
                weight_sum += wj;
                smoothed_velocity += v[*j] * wj;
            }

            let wi = mass / d[index];
            weight_sum += wi;
            smoothed_velocity += v[index] * wi;

            if weight_sum > 0.0 {
                smoothed_velocity /= weight_sum;
            }

            *vel = smoothed_velocity;
        });

        let mut factor = time_step_in_seconds * self._pseudo_viscosity_coefficient;
        factor = crate::math_utils::clamp(factor, 0.0, 1.0);

        (self._particle_system_data.write().unwrap().velocities_mut(),
         &mut smoothed_velocities).into_par_iter().for_each(|(vel, smooth_vel)| {
            *vel = *vel * (1.0 - factor) + *smooth_vel * factor;
        });
    }
}

/// Shared pointer type for the SphSolver3.
pub type SphSolver3Ptr = Arc<RwLock<SphSolver3>>;