qtruss 0.3.0

//! Describes a 2D truss.

use maria_linalg::{
    Matrix,
    Vector,
};

use super::Element;

#[derive(Clone, Copy, PartialEq, Debug)]
/// Enumerates the types of constraints available on
/// truss nodes.
pub enum Constraint {
    /// A free node
    Free (Vector<2>),
    /// A pinned node
    Pin,
    /// A node free to slide in the X direction, but not in the Y direction.
    HorizontalSlide (Vector<2>),
    /// A node free to slide in the Y direction, but not in the X direction.
    VerticalSlide (Vector<2>),
}

#[derive(Clone, Copy, PartialEq, Debug)]
/// A 2D truss node.
pub struct Node {
    pub location: Vector<2>,
    pub constraint: Constraint,
}

impl Node {
    /// Constructs an empty "zero" node.
    pub fn zero() -> Self {
        Self {
            location: Vector::zero(),
            constraint: Constraint::Free (Vector::zero()),
        }
    }

    /// Constructs a new node.
    pub fn new(
        location: Vector<2>,
        constraint: Constraint,
    ) -> Self {
        Self {
            location,
            constraint,
        }
    }
}

#[derive(Clone, Copy, Debug)]
/// A 2D truss with `N` nodes, `K` elements, and `F` degrees of freedom.
pub struct Truss<const N: usize, const K: usize, const F: usize> {
    nodes: [Node; N],
    pub elements: [Element; K],
    i: usize,
    pub forces: Option<Vector<F>>,
    pub displacements: Option<Vector<F>>,
}

impl<const N: usize, const K: usize, const F: usize> Truss<N, K, F> {
    pub fn new(nodes: [Node; N]) -> Self {
        Self {
            nodes,
            elements: [Element::zero(); K],
            i: 0,
            forces: None,
            displacements: None,
        }
    }

    /// Add an element to this truss and returns the index of this element.
    /// 
    /// Element indices begin at `0` and end at `K - 1`.
    pub fn add(&mut self, one: usize, two: usize, e: f64, a: f64) -> Option<usize> {
        // Quit if the node numbers are invalid
        if one >= N {
            println!("Invalid node number: {}", one);
            return None;
        }
        if two >= N {
            println!("Invalid node number: {}", two);
            return None;
        }

        let node_one = self.nodes[one];
        let node_two = self.nodes[two];
        self.elements[self.i] = Element::new(
            one,
            node_one.location,
            two,
            node_two.location,
            e,
            a,
        );

        let output = self.i;
        self.i += 1;

        Some (output)
    }

    /// Assemble the global stiffness matrix for this truss.
    /// 
    /// The global stiffness matrix has dimensions `(F, F)`.
    pub fn global_stiffness_matrix(&self) -> Option<Matrix<F>> {
        // If we don't have enough elements, quit
        if self.i < K {
            return None;
        }

        // Create a new `F` by `F` matrix
        let mut matrix = Matrix::zero();

        // Assemble a list of free degrees (up to two per node)
        // `degrees[i]` represents the global indices (if they exist) at node `i`
        let mut degrees: [(Option<usize>, Option<usize>); N] = [(None, None); N];
        let mut f = 0;
        for (n, node) in self.nodes.iter().enumerate() {
            // Create an X degree of freedom, if necessary
            if matches!(node.constraint, Constraint::Free (_))
                || matches!(node.constraint, Constraint::HorizontalSlide (_))
            {
                degrees[n].0 = Some (f);
                f += 1;
            }

            // Create a Y degree of freedom, if necessary
            if matches!(node.constraint, Constraint::Free (_))
                || matches!(node.constraint, Constraint::VerticalSlide (_))
            {
                degrees[n].1 = Some (f);
                f += 1;
            }
        }

        for element in self.elements {
            let attached: [Option<usize>; 4] = [
                degrees[element.one].0,
                degrees[element.one].1,
                degrees[element.two].0,
                degrees[element.two].1,
            ];

            for (i, a0) in attached.iter().enumerate() {
                for (j, a1) in attached.iter().enumerate() {
                    if let Some (arow) = a0 {
                        if let Some (acol) = a1 {
                            matrix[(*arow, *acol)] += element.stiffness[(i, j)];
                        }
                    }
                }
            }
        }

        Some (matrix)
    }

    /// Solves this truss.
    pub fn solve(&mut self) -> Option<Vector<F>> {
        // Get forces applied
        let mut f = Vector::zero();
        let mut i = 0;

        for node in self.nodes {
            match node.constraint {
                Constraint::Pin => (),
                Constraint::Free (applied) => {
                    f[i] = applied[0];
                    f[i + 1] = applied[1];
                    i += 2;
                }
                Constraint::HorizontalSlide (applied) => {
                    f[i] = applied[0];
                    i += 1;
                }
                Constraint::VerticalSlide (applied) => {
                    f[i] = applied[0];
                    i += 1;
                }
            }
        }

        // Store forces
        self.forces = Some (f);

        // Get global stiffness matrix
        let k = match self.global_stiffness_matrix() {
            Some (mx) => mx,
            None => {
                println!("Global stiffness matrix is singular.");
                return None;
            },
        };

        // Store displacements
        self.displacements = Some (k.inverse().mult(f));

        // Return displacements
        self.displacements
    }

    /// Gets the displacements of a provided node, given its index.
    /// 
    /// *Note*: you must solve this truss by calling `Truss::solve()` before calling this function.
    pub fn displacement(&self, node: usize) -> Option<Vector<2>> {
        // Get global displacements
        let displacements = match self.displacements {
            Some (d) => d,
            None => {
                println!("This truss has not yet been solved.");
                return None;
            },
        };

        // Start a counter
        let mut i = 0;

        // Iterate through the nodes, up to the current node index
        for k in 0..node {
            // Get the current node
            let n = self.nodes[k];

            i += match n.constraint {
                Constraint::Pin => 0,
                Constraint::Free (_) => 2,
                Constraint::HorizontalSlide (_) => 1,
                Constraint::VerticalSlide (_) => 1,
            };
        }

        let mut output: Vector<2> = Vector::zero();

        match self.nodes[node].constraint {
            // A pin doesn't move
            Constraint::Pin => return Some (output),
            
            // A free node can move in either dimension
            Constraint::Free (_) => {
                output[0] = displacements[i];
                output[1] = displacements[i + 1];
            },

            // A slider can move in only one dimension
            Constraint::HorizontalSlide (_) => output[0] = displacements[i],
            Constraint::VerticalSlide (_) => output[1] = displacements[i],
        }

        Some (output)
    }

    /// Computes the internal force in a provided element, given its index.
    /// 
    /// *Note*: you must solve this truss by calling `Truss::solve()` before calling this function.
    pub fn internal_force(&self, elt: usize) -> Option<f64> {
        // Get this element
        let element = self.elements[elt];

        // Get this element's displacements
        let mut u: Vector<4> = Vector::zero();

        let left_constraint = self.nodes[element.one].constraint;
        let right_constraint = self.nodes[element.two].constraint;

        let left_displacement = match self.displacement(element.one) {
            Some (v) => v,
            None => {
                println!("This truss has not yet been solved.");
                return None;
            },
        };
        match left_constraint {
            Constraint::Pin => {
                u[0] = 0.0;
                u[1] = 0.0;
            },
            Constraint::Free (_) => {
                u[0] = left_displacement[0];
                u[1] = left_displacement[1];
            },
            Constraint::HorizontalSlide (_) => {
                u[0] = left_displacement[0];
                u[1] = 0.0;
            },
            Constraint::VerticalSlide (_) => {
                u[0] = 0.0;
                u[1] = left_displacement[1];
            },
        }

        let right_displacement = match self.displacement(element.two) {
            Some (v) => v,
            None => {
                println!("This truss has not yet been solved.");
                return None;
            },
        };
        match right_constraint {
            Constraint::Pin => {
                u[2] = 0.0;
                u[3] = 0.0;
            },
            Constraint::Free (_) => {
                u[2] = right_displacement[0];
                u[3] = right_displacement[1];
            },
            Constraint::HorizontalSlide (_) => {
                u[2] = right_displacement[0];
                u[3] = 0.0;
            },
            Constraint::VerticalSlide (_) => {
                u[2] = 0.0;
                u[3] = right_displacement[1];
            },
        }
        
        // Get this element's stiffness matrix
        let k = element.stiffness;

        // Get the reaction forces on this element
        let force = k.mult(u);

        // Get the internal force in this element
        // 
        // Note: for reasons of convention, we always take the last two
        // values of the reaction forces vector
        let internal = Vector::new([
            force[2],
            force[3],
        ]);

        // Get the direction of this element
        let direction = element.direction;

        // Return the internal force of this element
        Some (internal.dot(direction))
    }

    /// Computes the internal stress in a provided element, in Pa.
    /// 
    /// *Note*: you must solve this truss by calling `Truss::solve()` before calling this function.
    pub fn stress(&self, elt: usize) -> Option<f64> {
        if let Some (f) = self.internal_force(elt) {
            let area = self.elements[elt].area;
            Some (f / area)
        } else {
            println!("This truss has not yet been solved.");
            None
        }
    }

    /// Computes the compliance of this truss in Joules.
    /// 
    /// *Note*: you must solve this truss by calling `Truss::solve()` before calling this function.
    pub fn compliance(&self) -> Option<f64> {
        let f = match self.forces {
            Some (v) => v,
            None => {
                println!("This truss has not yet been solved.");
                return None;
            },
        };

        let d = match self.displacements {
            Some (v) => v,
            None => {
                println!("This truss has not yet been solved.");
                return None;
            },
        };

        Some (f.dot(d))
    }

    /// Computes the total volume of material contained in this truss.
    pub fn volume(&self) -> f64 {
        let mut volume = 0.0;

        for element in self.elements {
            volume += element.area * element.length;
        }

        volume
    }
}