State.hpp

Go to the documentation of this file.

#pragma once
 
#include <polyfem/Common.hpp>
 
#include <polyfem/Units.hpp>
 
#include <polyfem/basis/ElementBases.hpp>
#include <polyfem/basis/InterfaceData.hpp>
 
#include <polyfem/assembler/ElementAssemblyValues.hpp>
#include <polyfem/assembler/AssemblyValsCache.hpp>
#include <polyfem/assembler/RhsAssembler.hpp>
#include <polyfem/assembler/PressureAssembler.hpp>
#include <polyfem/assembler/MacroStrain.hpp>
#include <polyfem/assembler/Problem.hpp>
#include <polyfem/assembler/Assembler.hpp>
#include <polyfem/assembler/AssemblerUtils.hpp>
 
#include <polyfem/mesh/Mesh.hpp>
#include <polyfem/mesh/Obstacle.hpp>
#include <polyfem/mesh/MeshNodes.hpp>
#include <polyfem/mesh/LocalBoundary.hpp>
 
#include <polyfem/solver/SolveData.hpp>
#include <polyfem/solver/DiffCache.hpp>
 
#include <polyfem/utils/StringUtils.hpp>
#include <polyfem/utils/ElasticityUtils.hpp>
#include <polyfem/utils/JSONUtils.hpp>
#include <polyfem/utils/Logger.hpp>
#include <polyfem/assembler/PeriodicBoundary.hpp>
 
#include <polyfem/io/OutData.hpp>
 
#include <polysolve/linear/Solver.hpp>
 
#include <Eigen/Dense>
#include <Eigen/Sparse>
 
#include <memory>
#include <string>
#include <unordered_map>
 
#include <spdlog/sinks/basic_file_sink.h>
 
#include <ipc/collision_mesh.hpp>
#include <ipc/utils/logger.hpp>
 
// Forward declaration
namespace polysolve::nonlinear
{
    class Solver;
}
 
namespace polyfem::assembler
{
    class Mass;
    class ViscousDamping;
    class ViscousDampingPrev;
} // namespace polyfem::assembler
 
namespace polyfem
{
    namespace mesh
    {
        class Mesh2D;
        class Mesh3D;
    } // namespace mesh
 
    enum class CacheLevel
    {
        None,
        Solution,
        Derivatives
    };
 
    class State
    {
    public:
        //---------------------------------------------------
        //-----------------initialization--------------------
        //---------------------------------------------------
 
        ~State() = default;
        State();
 
        void set_max_threads(const int max_threads = std::numeric_limits<int>::max());
 
        void init(const json &args, const bool strict_validation);
 
        void init_time();
 
        json args;
 
        //---------------------------------------------------
        //-----------------logger----------------------------
        //---------------------------------------------------
 
        void init_logger(
            const std::string &log_file,
            const spdlog::level::level_enum log_level,
            const spdlog::level::level_enum file_log_level,
            const bool is_quiet);
 
        void init_logger(std::ostream &os, const spdlog::level::level_enum log_level);
 
        void set_log_level(const spdlog::level::level_enum log_level);
 
        std::string get_log(const Eigen::MatrixXd &sol)
        {
            std::stringstream ss;
            save_json(sol, ss);
            return ss.str();
        }
 
    private:
        void init_logger(const std::vector<spdlog::sink_ptr> &sinks, const spdlog::level::level_enum log_level);
 
        spdlog::sink_ptr console_sink_ = nullptr;
        spdlog::sink_ptr file_sink_ = nullptr;
 
    public:
        //---------------------------------------------------
        //-----------------assembly--------------------------
        //---------------------------------------------------
 
        Units units;
 
 
        std::shared_ptr<assembler::Assembler> assembler = nullptr;
 
        std::shared_ptr<assembler::Mass> mass_matrix_assembler = nullptr;
 
        std::shared_ptr<assembler::MixedAssembler> mixed_assembler = nullptr;
        std::shared_ptr<assembler::Assembler> pressure_assembler = nullptr;
 
        std::shared_ptr<assembler::PressureAssembler> elasticity_pressure_assembler = nullptr;
 
        std::shared_ptr<assembler::ViscousDamping> damping_assembler = nullptr;
        std::shared_ptr<assembler::ViscousDampingPrev> damping_prev_assembler = nullptr;
 
        std::shared_ptr<assembler::Problem> problem;
 
        std::vector<basis::ElementBases> bases;
        std::vector<basis::ElementBases> pressure_bases;
        std::vector<basis::ElementBases> geom_bases_;
 
        int n_bases;
        int n_pressure_bases;
        int n_geom_bases;
 
        std::map<int, Eigen::MatrixXd> polys;
        std::map<int, std::pair<Eigen::MatrixXd, Eigen::MatrixXi>> polys_3d;
 
        Eigen::VectorXi disc_orders;
 
        std::shared_ptr<polyfem::mesh::MeshNodes> mesh_nodes, geom_mesh_nodes, pressure_mesh_nodes;
 
        assembler::AssemblyValsCache ass_vals_cache;
        assembler::AssemblyValsCache mass_ass_vals_cache;
        assembler::AssemblyValsCache pressure_ass_vals_cache;
 
        StiffnessMatrix mass;
        double avg_mass;
 
        Eigen::MatrixXd rhs;
 
        bool use_avg_pressure;
 
        std::string formulation() const;
 
        bool iso_parametric() const;
 
        const std::vector<basis::ElementBases> &geom_bases() const
        {
            return iso_parametric() ? bases : geom_bases_;
        }
 
        void build_basis();
        void assemble_rhs();
        void assemble_mass_mat();
 
        std::shared_ptr<assembler::RhsAssembler> build_rhs_assembler(
            const int n_bases,
            const std::vector<basis::ElementBases> &bases,
            const assembler::AssemblyValsCache &ass_vals_cache) const;
        std::shared_ptr<assembler::RhsAssembler> build_rhs_assembler() const
        {
            return build_rhs_assembler(n_bases, bases, mass_ass_vals_cache);
        }
 
        std::shared_ptr<assembler::PressureAssembler> build_pressure_assembler(
            const int n_bases_,
            const std::vector<basis::ElementBases> &bases_) const;
        std::shared_ptr<assembler::PressureAssembler> build_pressure_assembler() const
        {
            return build_pressure_assembler(n_bases, bases);
        }
 
        int n_boundary_samples() const
        {
            using assembler::AssemblerUtils;
            const int n_b_samples_j = args["space"]["advanced"]["n_boundary_samples"];
            const int gdiscr_order = mesh->orders().size() <= 0 ? 1 : mesh->orders().maxCoeff();
            const int discr_order = std::max(disc_orders.maxCoeff(), gdiscr_order);
 
            const int n_b_samples = std::max(n_b_samples_j, AssemblerUtils::quadrature_order("Mass", discr_order, AssemblerUtils::BasisType::POLY, mesh->dimension()));
 
            return n_b_samples;
        }
 
    private:
        void sol_to_pressure(Eigen::MatrixXd &sol, Eigen::MatrixXd &pressure);
        void build_polygonal_basis();
 
    public:
        void set_materials(std::vector<std::shared_ptr<assembler::Assembler>> &assemblers) const;
        void set_materials(assembler::Assembler &assembler) const;
 
        //---------------------------------------------------
        //-----------------solver----------------------------
        //---------------------------------------------------
 
    public:
        void solve_problem(Eigen::MatrixXd &sol, Eigen::MatrixXd &pressure);
        void solve(Eigen::MatrixXd &sol, Eigen::MatrixXd &pressure)
        {
            if (!mesh)
            {
                logger().error("Load the mesh first!");
                return;
            }
            stats.compute_mesh_stats(*mesh);
 
            build_basis();
 
            assemble_rhs();
            assemble_mass_mat();
 
            solve_export_to_file = false;
            solution_frames.clear();
            solve_problem(sol, pressure);
            solve_export_to_file = true;
        }
 
        solver::SolveData solve_data;
        void init_solve(Eigen::MatrixXd &sol, Eigen::MatrixXd &pressure);
        void solve_transient_navier_stokes_split(const int time_steps, const double dt, Eigen::MatrixXd &sol, Eigen::MatrixXd &pressure);
        void solve_transient_navier_stokes(const int time_steps, const double t0, const double dt, Eigen::MatrixXd &sol, Eigen::MatrixXd &pressure);
        void solve_transient_linear(const int time_steps, const double t0, const double dt, Eigen::MatrixXd &sol, Eigen::MatrixXd &pressure);
        void solve_transient_tensor_nonlinear(const int time_steps, const double t0, const double dt, Eigen::MatrixXd &sol);
        void init_nonlinear_tensor_solve(Eigen::MatrixXd &sol, const double t = 1.0, const bool init_time_integrator = true);
        void init_linear_solve(Eigen::MatrixXd &sol, const double t = 1.0);
        void initial_solution(Eigen::MatrixXd &solution) const;
        void initial_velocity(Eigen::MatrixXd &velocity) const;
        void initial_acceleration(Eigen::MatrixXd &acceleration) const;
        void solve_linear(Eigen::MatrixXd &sol, Eigen::MatrixXd &pressure);
        void solve_navier_stokes(Eigen::MatrixXd &sol, Eigen::MatrixXd &pressure);
        void solve_tensor_nonlinear(Eigen::MatrixXd &sol, const int t = 0, const bool init_lagging = true);
 
        std::shared_ptr<polysolve::nonlinear::Solver> make_nl_solver(bool for_al) const;
 
        std::shared_ptr<utils::PeriodicBoundary> periodic_bc;
        bool has_periodic_bc() const
        {
            return args["boundary_conditions"]["periodic_boundary"]["enabled"].get<bool>();
        }
 
        void solve_linear(
            const std::unique_ptr<polysolve::linear::Solver> &solver,
            StiffnessMatrix &A,
            Eigen::VectorXd &b,
            const bool compute_spectrum,
            Eigen::MatrixXd &sol, Eigen::MatrixXd &pressure);
 
        bool is_problem_linear() const { return assembler->is_linear() && !is_contact_enabled() && !is_pressure_enabled(); }
 
    public:
        void build_stiffness_mat(StiffnessMatrix &stiffness);
 
        //---------------------------------------------------
        //-----------------nodes flags-----------------------
        //---------------------------------------------------
 
    public:
        std::unordered_map<int, std::array<bool, 3>>
        boundary_conditions_ids(const std::string &bc_type) const;
 
        std::vector<int> boundary_nodes;
        std::vector<int> pressure_boundary_nodes;
        std::vector<mesh::LocalBoundary> total_local_boundary;
        std::vector<mesh::LocalBoundary> local_boundary;
        std::vector<mesh::LocalBoundary> local_neumann_boundary;
        std::vector<mesh::LocalBoundary> local_pressure_boundary;
        std::unordered_map<int, std::vector<mesh::LocalBoundary>> local_pressure_cavity;
        std::map<int, basis::InterfaceData> poly_edge_to_data;
        std::vector<int> dirichlet_nodes;
        std::vector<RowVectorNd> dirichlet_nodes_position;
        std::vector<int> neumann_nodes;
        std::vector<RowVectorNd> neumann_nodes_position;
 
        Eigen::VectorXi in_node_to_node;
        Eigen::VectorXi in_primitive_to_primitive;
 
        std::vector<int> primitive_to_node() const;
        std::vector<int> node_to_primitive() const;
 
    private:
        void build_node_mapping();
 
        //---------------------------------------------------
        //-----------------Geometry--------------------------
        //---------------------------------------------------
    public:
        std::unique_ptr<mesh::Mesh> mesh;
        mesh::Obstacle obstacle;
 
        void load_mesh(bool non_conforming = false,
                       const std::vector<std::string> &names = std::vector<std::string>(),
                       const std::vector<Eigen::MatrixXi> &cells = std::vector<Eigen::MatrixXi>(),
                       const std::vector<Eigen::MatrixXd> &vertices = std::vector<Eigen::MatrixXd>());
 
        void load_mesh(GEO::Mesh &meshin, const std::function<int(const size_t, const std::vector<int> &, const RowVectorNd &, bool)> &boundary_marker, bool non_conforming = false, bool skip_boundary_sideset = false);
 
        void load_mesh(const Eigen::MatrixXd &V, const Eigen::MatrixXi &F, bool non_conforming = false)
        {
            mesh = mesh::Mesh::create(V, F, non_conforming);
            load_mesh(non_conforming);
        }
 
        void reset_mesh();
 
        void build_mesh_matrices(Eigen::MatrixXd &V, Eigen::MatrixXi &F);
 
        bool remesh(const double time, const double dt, Eigen::MatrixXd &sol);
 
        //---------------------------------------------------
        //-----------------IPC-------------------------------
        //---------------------------------------------------
 
        ipc::CollisionMesh collision_mesh;
 
        ipc::CollisionMesh periodic_collision_mesh;
        Eigen::VectorXi periodic_collision_mesh_to_basis;
 
        static void build_collision_mesh(
            const mesh::Mesh &mesh,
            const int n_bases,
            const std::vector<basis::ElementBases> &bases,
            const std::vector<basis::ElementBases> &geom_bases,
            const std::vector<mesh::LocalBoundary> &total_local_boundary,
            const mesh::Obstacle &obstacle,
            const json &args,
            const std::function<std::string(const std::string &)> &resolve_input_path,
            const Eigen::VectorXi &in_node_to_node,
            ipc::CollisionMesh &collision_mesh);
 
        void build_collision_mesh();
        void build_periodic_collision_mesh();
 
        bool is_obstacle_vertex(const size_t vi) const
        {
            // The obstalce vertices are at the bottom of the collision mesh vertices
            return vi >= collision_mesh.full_num_vertices() - obstacle.n_vertices();
        }
 
        bool is_contact_enabled() const
        {
            return args["contact"]["enabled"];
        }
 
        bool is_pressure_enabled() const
        {
            return (args["boundary_conditions"]["pressure_boundary"].size() > 0)
                   || (args["boundary_conditions"]["pressure_cavity"].size() > 0);
        }
 
        bool has_dhat = false;
 
        //---------------------------------------------------
        //-----------------OUTPUT----------------------------
        //---------------------------------------------------
    public:
        std::string output_dir;
 
        bool solve_export_to_file = true;
        std::vector<io::SolutionFrame> solution_frames;
        io::OutGeometryData out_geom;
        io::OutRuntimeData timings;
        io::OutStatsData stats;
        double starting_min_edge_length = -1;
        double starting_max_edge_length = -1;
        double min_boundary_edge_length = -1;
 
        void export_data(const Eigen::MatrixXd &sol, const Eigen::MatrixXd &pressure);
 
        void save_timestep(const double time, const int t, const double t0, const double dt, const Eigen::MatrixXd &sol, const Eigen::MatrixXd &pressure);
 
        void save_subsolve(const int i, const int t, const Eigen::MatrixXd &sol, const Eigen::MatrixXd &pressure);
 
        void save_json(const Eigen::MatrixXd &sol, std::ostream &out);
 
        void save_json(const Eigen::MatrixXd &sol);
 
        void compute_errors(const Eigen::MatrixXd &sol);
 
        void save_restart_json(const double t0, const double dt, const int t) const;
 
        //-----------PATH management
        std::string root_path() const;
 
        std::string resolve_input_path(const std::string &path, const bool only_if_exists = false) const;
 
        std::string resolve_output_path(const std::string &path) const;
 
        //---------------------------------------------------
        //-----------------differentiable--------------------
        //---------------------------------------------------
    public:
        solver::CacheLevel optimization_enabled = solver::CacheLevel::None;
        void cache_transient_adjoint_quantities(const int current_step, const Eigen::MatrixXd &sol, const Eigen::MatrixXd &disp_grad);
        solver::DiffCache diff_cached;
 
        std::unique_ptr<polysolve::linear::Solver> lin_solver_cached; // matrix factorization of last linear solve
 
        int ndof() const
        {
            const int actual_dim = problem->is_scalar() ? 1 : mesh->dimension();
            if (mixed_assembler == nullptr)
                return actual_dim * n_bases;
            else
                return actual_dim * n_bases + n_pressure_bases;
        }
 
        // Aux functions for setting up adjoint equations
        void compute_force_jacobian(const Eigen::MatrixXd &sol, const Eigen::MatrixXd &disp_grad, StiffnessMatrix &hessian);
        void compute_force_jacobian_prev(const int force_step, const int sol_step, StiffnessMatrix &hessian_prev) const;
        // Solves the adjoint PDE for derivatives and caches
        void solve_adjoint_cached(const Eigen::MatrixXd &rhs);
        Eigen::MatrixXd solve_adjoint(const Eigen::MatrixXd &rhs) const;
        // Returns cached adjoint solve
        Eigen::MatrixXd get_adjoint_mat(int type) const
        {
            assert(diff_cached.adjoint_mat().size() > 0);
            if (problem->is_time_dependent())
            {
                if (type == 0)
                    return diff_cached.adjoint_mat().leftCols(diff_cached.adjoint_mat().cols() / 2);
                else if (type == 1)
                    return diff_cached.adjoint_mat().middleCols(diff_cached.adjoint_mat().cols() / 2, diff_cached.adjoint_mat().cols() / 2);
                else
                    log_and_throw_adjoint_error("Invalid adjoint type!");
            }
 
            return diff_cached.adjoint_mat();
        }
        Eigen::MatrixXd solve_static_adjoint(const Eigen::MatrixXd &adjoint_rhs) const;
        Eigen::MatrixXd solve_transient_adjoint(const Eigen::MatrixXd &adjoint_rhs) const;
        // Change geometric node positions
        void set_mesh_vertex(int v_id, const Eigen::VectorXd &vertex);
        void get_vertices(Eigen::MatrixXd &vertices) const;
        void get_elements(Eigen::MatrixXi &elements) const;
 
        // Get geometric node indices for surface/volume
        void compute_surface_node_ids(const int surface_selection, std::vector<int> &node_ids) const;
        void compute_total_surface_node_ids(std::vector<int> &node_ids) const;
        void compute_volume_node_ids(const int volume_selection, std::vector<int> &node_ids) const;
 
        // to replace the initial condition in json during initial condition optimization
        Eigen::MatrixXd initial_sol_update, initial_vel_update;
        // mapping from positions of FE basis nodes to positions of geometry nodes
        StiffnessMatrix basis_nodes_to_gbasis_nodes;
 
        //---------------------------------------------------
        //-----------------homogenization--------------------
        //---------------------------------------------------
    public:
        assembler::MacroStrainValue macro_strain_constraint;
 
        void solve_homogenization_step(Eigen::MatrixXd &sol, const int t = 0, bool adaptive_initial_weight = false); // sol is the extended solution, i.e. [periodic fluctuation, macro strain]
        void init_homogenization_solve(const double t);
        void solve_homogenization(const int time_steps, const double t0, const double dt, Eigen::MatrixXd &sol);
        bool is_homogenization() const
        {
            return args["boundary_conditions"]["periodic_boundary"]["linear_displacement_offset"].size() > 0;
        }
    };
 
} // namespace polyfem