epanet-rs 0.2.2

//! `Valve` link and `ValveType` variants (PRV, PSV, PBV, FCV, TCV, GPV) with their control logic.

use crate::constants::*;
use crate::model::curve::{Curve, ValveCurve};
use crate::model::link::{LinkCoefficients, LinkStatus, LinkTrait, NodeModification};
use crate::model::options::SimulationOptions;
use crate::model::units::{Ft, UnitConversion, UnitSystem};
use serde::{Deserialize, Serialize};

#[derive(Debug, PartialEq, Eq, Deserialize, Serialize, Clone, strum::Display)]
pub enum ValveType {
    PRV, // Pressure Reducing Valve
    PSV, // Pressure Sensing Valve
    PBV, // Pressure Breaking Valve
    FCV, // Flow Control Valve
    TCV, // Throttle Control Valve
    PCV, // Positional Control Valve
    GPV, // General Purpose Valve
}

#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct Valve {
    pub diameter: Ft,
    pub setting: f64,
    pub curve_id: Option<Box<str>>,
    pub valve_type: ValveType,
    pub minor_loss: f64,
    #[serde(skip)]
    pub gpv_curve: Option<ValveCurve>,
    #[serde(skip)]
    pub pcv_curve: Option<Curve>,
}

impl LinkTrait for Valve {
    #[inline]
    fn coefficients(
        &self,
        q: f64,
        _resistance: f64,
        setting: f64,
        status: LinkStatus,
        excess_flow_upstream: f64,
        excess_flow_downstream: f64,
        elevation_upstream: f64,
        elevation_downstream: f64,
    ) -> LinkCoefficients {
        // if the valve is closed, fixed closed, or XPressure, return a high resistance valve
        if status == LinkStatus::Closed
            || status == LinkStatus::TempClosed
            || status == LinkStatus::FixedClosed
        {
            return LinkCoefficients::simple(1.0 / BIG_VALUE, q);
        }

        // if the valve is open or fixed open, use the minor loss coefficient to compute the coefficients
        if status == LinkStatus::Open || status == LinkStatus::FixedOpen {
            let km = 0.02517 * self.minor_loss / self.diameter.powi(4);
            let (g_inv, y) = self.valve_coefficients(q, km);
            return LinkCoefficients::simple(g_inv, y);
        }

        match self.valve_type {
            // Throttle Control Valve (TCV)
            ValveType::TCV => {
                // Minor loss coefficient is the setting of the valve
                let km = 0.02517 * setting / self.diameter.powi(4);
                let (g_inv, y) = self.valve_coefficients(q, km);
                LinkCoefficients::simple(g_inv, y)
            }
            // Pressure Breaking Valve (PBV)
            ValveType::PBV => LinkCoefficients::simple(BIG_VALUE, setting * BIG_VALUE),
            // Positional Control Valve (PCV)
            ValveType::PCV => {
                let km = self.pcv_minor_loss(setting);
                let (g_inv, y) = self.valve_coefficients(q, km);
                LinkCoefficients::simple(g_inv, y)
            }
            // Flow Control Valve (FCV)
            ValveType::FCV => {
                let (g_inv, y) = self.valve_coefficients(q, SMALL_VALUE);
                // if flow is less than the setting, treat as a regular valve (no flow control/restrictions)
                if q < setting {
                    LinkCoefficients::simple(g_inv, y)
                } else {
                    let hloss = y / g_inv + BIG_VALUE * (q - setting);
                    let hgrad = BIG_VALUE;
                    LinkCoefficients::simple(1.0 / hgrad, hloss / hgrad)
                }
            }
            // Pressure Reducing Valve (PRV) — setting is pressure (in feet of head)
            // at the downstream node; convert to absolute head.
            ValveType::PRV => {
                if status == LinkStatus::Active {
                    let target_head = setting + elevation_downstream;
                    self.prv_coefficients(target_head, excess_flow_downstream)
                } else {
                    LinkCoefficients::simple(1.0 / SMALL_VALUE, q)
                }
            }
            // Pressure Sustaining Valve (PSV) — setting is pressure (in feet of head)
            // at the upstream node; convert to absolute head.
            ValveType::PSV => {
                if status == LinkStatus::Active {
                    let target_head = setting + elevation_upstream;
                    self.psv_coefficients(target_head, excess_flow_upstream)
                } else {
                    LinkCoefficients::simple(1.0 / SMALL_VALUE, q)
                }
            }
            ValveType::GPV => self.gpv_coefficients(q),
        }
    }
    /// Return the resistance of the valve
    fn resistance(&self) -> f64 {
        SMALL_VALUE
    }

    fn update_status(
        &self,
        setting: f64,
        status: LinkStatus,
        q: f64,
        head_upstream: f64,
        head_downstream: f64,
        elevation_upstream: f64,
        elevation_downstream: f64,
    ) -> Option<LinkStatus> {
        match self.valve_type {
            ValveType::PRV => {
                let target_head = setting + elevation_downstream;
                self.prv_status(target_head, status, q, head_upstream, head_downstream)
            }
            ValveType::PSV => {
                let target_head = setting + elevation_upstream;
                self.psv_status(target_head, status, q, head_upstream, head_downstream)
            }
            _ => None,
        }
    }

    fn initial_flow(&self) -> f64 {
        // return the initial flow of the valve corresponding to 1 ft/s velocity
        0.25 * std::f64::consts::PI * self.diameter.powi(2)
    }
}
impl Valve {
    /// Compute the updated status of the pressure reducing valve
    fn prv_status(
        &self,
        setting: f64,
        status: LinkStatus,
        q: f64,
        head_upstream: f64,
        head_downstream: f64,
    ) -> Option<LinkStatus> {
        match status {
            LinkStatus::Active => {
                if q < -Q_TOL {
                    Some(LinkStatus::Closed)
                }
                // closed if flow is negative
                else if head_upstream < setting - H_TOL {
                    Some(LinkStatus::Open)
                }
                // open if head upstream is less than the setting
                else {
                    None
                } // no status change
            }
            LinkStatus::Open => {
                if q < -Q_TOL {
                    Some(LinkStatus::Closed)
                }
                // closed if flow is negative
                else if head_downstream > setting + H_TOL {
                    Some(LinkStatus::Active)
                }
                // active if head downstream is greater than the setting
                else {
                    None
                } // no status change
            }
            LinkStatus::Closed => {
                if head_upstream >= setting + H_TOL && head_downstream < setting - H_TOL {
                    Some(LinkStatus::Active) // active if head upstream is greater than the setting and head downstream is less than the setting
                } else if head_upstream < setting - H_TOL && head_upstream > head_downstream + H_TOL
                {
                    Some(LinkStatus::Open) // open if head upstream is less than the setting and head upstream is greater than head downstream plus the tolerance
                } else {
                    None // no status change
                }
            }
            LinkStatus::XPressure => {
                if q < -Q_TOL {
                    Some(LinkStatus::Closed)
                } else {
                    None
                }
            }
            _ => None,
        }
    }

    /// Compute the updated status of the pressure sustaining valve
    fn psv_status(
        &self,
        setting: f64,
        status: LinkStatus,
        q: f64,
        head_upstream: f64,
        head_downstream: f64,
    ) -> Option<LinkStatus> {
        match status {
            LinkStatus::Active => {
                if q < -Q_TOL {
                    Some(LinkStatus::Closed)
                }
                // closed if flow is negative
                else if head_downstream > setting + H_TOL {
                    Some(LinkStatus::Open)
                }
                //  open if head downstream is less then the setting
                else {
                    None
                } // no status change
            }
            LinkStatus::Open => {
                if q < -Q_TOL {
                    Some(LinkStatus::Closed)
                }
                // closed if flow is negative
                else if head_upstream < setting - H_TOL {
                    Some(LinkStatus::Active)
                }
                // Active when upstream head is less than the setting
                else {
                    None
                } // no status change
            }
            LinkStatus::Closed => {
                if head_downstream > setting + H_TOL && head_upstream > head_downstream + H_TOL {
                    Some(LinkStatus::Open)
                } else if head_upstream > setting + H_TOL && head_upstream > head_downstream + H_TOL
                {
                    Some(LinkStatus::Active)
                } else {
                    None
                } // no status change
            }
            LinkStatus::XPressure => {
                if q < -Q_TOL {
                    Some(LinkStatus::Closed)
                } else {
                    None
                }
            }
            _ => None,
        }
    }

    /// Compute the coefficients for general purpose valve with a flow q
    fn gpv_coefficients(&self, q: f64) -> LinkCoefficients {
        let curve = self.gpv_curve.as_ref().unwrap();

        let q_abs = q.abs().max(TINY);

        let (h0, r) = curve.coefficients(q_abs);
        let r = r.max(TINY);

        LinkCoefficients::simple(1.0 / r, (h0 / r + q_abs) * q.signum())
    }

    /// Compute the coefficients for pressure sustaining valve with a flow q and excess flow upstream
    fn psv_coefficients(&self, setting: f64, excess_flow_upstream: f64) -> LinkCoefficients {
        let mut rhs_add = setting * BIG_VALUE;

        if excess_flow_upstream > 0.0 {
            rhs_add += excess_flow_upstream;
        }

        LinkCoefficients {
            g_inv: 0.0,
            y: -excess_flow_upstream,
            new_status: None,
            upstream_modification: Some(NodeModification {
                diagonal_add: BIG_VALUE,
                rhs_add,
            }),
            downstream_modification: None,
        }
    }

    /// Compute the coefficients for pressure reducing valve with a flow q and excess flow downstream
    fn prv_coefficients(&self, setting: f64, excess_flow_downstream: f64) -> LinkCoefficients {
        let mut rhs_add = setting * BIG_VALUE;
        if excess_flow_downstream < 0.0 {
            rhs_add += excess_flow_downstream;
        }

        LinkCoefficients {
            g_inv: 0.0,
            y: excess_flow_downstream,
            new_status: None,
            upstream_modification: None,
            downstream_modification: Some(NodeModification {
                diagonal_add: BIG_VALUE,
                rhs_add,
            }),
        }
    }

    fn pcv_minor_loss(&self, setting: f64) -> f64 {
        // Minor loss coefficient for a completely open valve
        let k_open = 0.02517 * self.minor_loss / self.diameter.powi(4);

        // Valve is completely closed
        if setting <= 0.0 {
            return BIG_VALUE;
        }
        // Valve is completely open
        if setting >= 100.0 {
            return k_open;
        }

        // Valve is partially open
        let ratio = if let Some(curve) = &self.pcv_curve {
            // Use the valve curve to compute the ratio
            let (h0, r) = curve.coefficients(setting);
            let ratio = h0 + r * setting;
            ratio / 100.0
        } else {
            // Assume a linear relationship between the setting and the minor loss coefficient
            setting / 100.0
        };

        // clamp the ratio to SMALL_VALUE and 1.0 to avoid division by zero
        let ratio = ratio.clamp(SMALL_VALUE, 1.0);

        // convert the ratio to a minor loss coefficient
        let km = k_open / ratio.powi(2);

        km.min(BIG_VALUE)
    }
    /// Compute the coefficients for throttle control valve with a minor loss coefficient km and flow q
    fn valve_coefficients(&self, q: f64, km: f64) -> (f64, f64) {
        if km > 0.0 {
            let q_abs = q.abs();
            let hgrad = 2.0 * km * q_abs;

            // guard against too small a head loss gradient
            if hgrad < RQ_TOL {
                let hgrad = RQ_TOL;
                let hloss = q * hgrad;
                (1.0 / hgrad, hloss / hgrad)
            } else {
                let hloss = q * hgrad / 2.0;
                (1.0 / hgrad, hloss / hgrad)
            }
        }
        // if no minor loss coefficient, use a low resistance linear head loss relation
        else {
            (1.0 / SMALL_VALUE, q)
        }
    }
}

impl UnitConversion for Valve {
    fn convert_to_standard(&mut self, options: &SimulationOptions) {
        if options.unit_system == UnitSystem::US {
            self.diameter /= 12.0; // convert in to ft
            // convert valve setting from PSI to feet
            if self.valve_type == ValveType::PRV
                || self.valve_type == ValveType::PSV
                || self.valve_type == ValveType::PBV
            {
                self.setting /= PSIperFT; // convert PSI to feet
            }
        } else {
            self.diameter /= 1000.0; // convert mm to m
        }

        // convert the diameter from the given unit system to feet
        self.diameter /= options.unit_system.per_feet();
        // convert valve setting from the given unit system to cfs
        if self.valve_type == ValveType::FCV {
            self.setting /= options.flow_units.per_cfs();
        }
        // convert valve setting from the given unit system to feet
        if self.valve_type == ValveType::PRV
            || self.valve_type == ValveType::PSV
            || self.valve_type == ValveType::PBV
        {
            self.setting /= options.unit_system.per_feet();
        }
    }
    fn convert_from_standard(&mut self, options: &SimulationOptions) {
        if options.unit_system == UnitSystem::US {
            self.diameter *= 12.0; // convert ft to in
            if self.valve_type == ValveType::PRV
                || self.valve_type == ValveType::PSV
                || self.valve_type == ValveType::PBV
            {
                self.setting *= PSIperFT; // convert feet to PSI
            }
        } else {
            self.diameter *= 1000.0; // convert m to mm
        }
        // convert the diameter from the given unit system to feet
        self.diameter *= options.unit_system.per_feet();
        // convert valve setting from the given unit system to cfs
        if self.valve_type == ValveType::FCV {
            self.setting *= options.flow_units.per_cfs();
        }
        // convert valve setting from the given unit system to feet
        if self.valve_type == ValveType::PRV
            || self.valve_type == ValveType::PSV
            || self.valve_type == ValveType::PBV
        {
            self.setting *= options.unit_system.per_feet();
        }
    }
}