SolveData.cpp

Go to the documentation of this file.

#include "SolveData.hpp"
 
#include <polyfem/solver/NLProblem.hpp>
#include <polyfem/solver/forms/Form.hpp>
#include <polyfem/solver/forms/lagrangian/BCLagrangianForm.hpp>
#include <polyfem/solver/forms/lagrangian/MatrixLagrangianForm.hpp>
#include <polyfem/solver/forms/lagrangian/PeriodicBoundaryLagrangianForm.hpp>
#include <polyfem/solver/forms/lagrangian/MacroStrainLagrangianForm.hpp>
#include <polyfem/solver/forms/BodyForm.hpp>
#include <polyfem/solver/forms/BarrierContactForm.hpp>
#include <polyfem/solver/forms/SmoothContactForm.hpp>
#include <polyfem/solver/forms/PressureForm.hpp>
#include <polyfem/solver/forms/PeriodicContactForm.hpp>
#include <polyfem/solver/forms/ElasticForm.hpp>
#include <polyfem/solver/forms/FrictionForm.hpp>
#include <polyfem/solver/forms/NormalAdhesionForm.hpp>
#include <polyfem/solver/forms/TangentialAdhesionForm.hpp>
#include <polyfem/solver/forms/InertiaForm.hpp>
#include <polyfem/solver/forms/LaggedRegForm.hpp>
#include <polyfem/solver/forms/RayleighDampingForm.hpp>
#include <polyfem/solver/forms/QuadraticPenaltyForm.hpp>
#include <polyfem/time_integrator/ImplicitTimeIntegrator.hpp>
#include <polyfem/assembler/ViscousDamping.hpp>
#include <polyfem/assembler/PressureAssembler.hpp>
#include <polyfem/assembler/Mass.hpp>
#include <polyfem/assembler/MacroStrain.hpp>
#include <polyfem/utils/Logger.hpp>
 
namespace polyfem::solver
{
    using namespace polyfem::time_integrator;
 
    namespace
    {
        std::string constraint_child(const std::string &group, const std::string &name)
        {
            return group.empty() || group == "/" ? "/" + name : group + "/" + name;
        }
 
        template <typename Callback>
        void with_constraint_resource(
            const io::ResourceIO &resources,
            const std::string &path,
            Callback &&callback)
        {
            if (resources.is_group(path))
            {
                callback(resources, path);
                return;
            }
 
            const io::HDF5IO constraint(resources.materialize(path));
            callback(constraint, "/");
        }
    } // namespace
 
    std::vector<std::shared_ptr<Form>> SolveData::init_forms(
        // General
        const Units &units,
        const io::ResourceIO &resources,
        const int dim,
        const double t,
        const Eigen::VectorXi &in_node_to_node,
 
        // Elastic form
        const int n_bases,
        std::vector<basis::ElementBases> &bases,
        const std::vector<basis::ElementBases> &geom_bases,
        const assembler::Assembler &assembler,
        assembler::AssemblyValsCache &ass_vals_cache,
        const assembler::AssemblyValsCache &mass_ass_vals_cache,
        const double jacobian_threshold,
        const ElementInversionCheck check_inversion,
        const unsigned conservative_max_iter,
 
        // Body form
        const int n_pressure_bases,
        const std::vector<int> &boundary_nodes,
        const std::vector<mesh::LocalBoundary> &local_boundary,
        const std::vector<mesh::LocalBoundary> &local_neumann_boundary,
        const QuadratureOrders &n_boundary_samples,
        const Eigen::MatrixXd &rhs,
        const Eigen::MatrixXd &sol,
        const assembler::Density &density,
 
        // Pressure form
        const std::vector<mesh::LocalBoundary> &local_pressure_boundary,
        const std::unordered_map<int, std::vector<mesh::LocalBoundary>> &local_pressure_cavity,
        const std::shared_ptr<assembler::PressureAssembler> pressure_assembler,
 
        // Inertia form
        const bool ignore_inertia,
        const StiffnessMatrix &mass,
        const std::shared_ptr<assembler::ViscousDamping> damping_assembler,
 
        // Lagged regularization form
        const double lagged_regularization_weight,
        const int lagged_regularization_iterations,
 
        // Augemented lagrangian form
        const size_t obstacle_ndof,
        const std::vector<std::string> &hard_constraint_files,
        const std::vector<json> &soft_constraint_files,
        const json &zero_mean,
 
        // Contact form
        const bool contact_enabled,
        const ipc::CollisionMesh &collision_mesh,
        const double dhat,
        const double avg_mass,
        const bool use_area_weighting,
        const bool use_improved_max_operator,
        const bool use_physical_barrier,
        const json &barrier_stiffness,
        const double initial_barrier_stiffness,
        const ipc::BroadPhaseMethod broad_phase,
        const double ccd_tolerance,
        const long ccd_max_iterations,
        const bool enable_shape_derivatives,
 
        // Smooth Contact Form
        const bool use_gcp_formulation,
        const double alpha_t,
        const double alpha_n,
        const bool use_adaptive_dhat,
        const double min_distance_ratio,
 
        // Normal Adhesion Form
        const bool adhesion_enabled,
        const double dhat_p,
        const double dhat_a,
        const double Y,
 
        // Tangential Adhesion Form
        const double tangential_adhesion_coefficient,
        const double epsa,
        const int tangential_adhesion_iterations,
 
        // Homogenization
        const assembler::MacroStrainValue &macro_strain_constraint,
 
        // Periodic contact
        const bool periodic_contact,
        const Eigen::VectorXi &tiled_to_single,
 
        // Friction form
        const double friction_coefficient,
        const double epsv,
        const int friction_iterations,
 
        // Rayleigh damping form
        const json &rayleigh_damping,
 
        // BC augmented-Lagrangian mass-metric lumping
        const BCLumpingMode al_lumping,
 
        // Boundary-ID periodic constraints
        const mesh::Mesh *periodic_mesh,
        const std::vector<mesh::LocalBoundary> *periodic_local_boundary,
        const json &periodic_conditions,
        const int fe_space_id)
    {
        const bool is_time_dependent = time_integrator != nullptr;
        assert(!is_time_dependent || time_integrator != nullptr);
        const double dt = is_time_dependent ? time_integrator->dt() : 0.0;
        const int ndof = n_bases * dim;
        // if (is_formulation_mixed) // mixed not supported
        //  ndof_ += n_pressure_bases; // pressure is a scalar
        const bool is_volume = dim == 3;
 
        std::vector<std::shared_ptr<Form>> forms;
        al_form.clear();
 
        elastic_form = std::make_shared<ElasticForm>(
            n_bases, bases, geom_bases, assembler, ass_vals_cache,
            t, dt, is_volume, jacobian_threshold, check_inversion, conservative_max_iter);
        forms.push_back(elastic_form);
 
        if (rhs_assembler != nullptr)
        {
            body_form = std::make_shared<BodyForm>(
                ndof, n_pressure_bases, boundary_nodes, local_boundary,
                local_neumann_boundary, n_boundary_samples, rhs, *rhs_assembler,
                density, /*is_formulation_mixed=*/false,
                is_time_dependent);
            body_form->update_quantities(t, sol);
            forms.push_back(body_form);
        }
 
        if (pressure_assembler != nullptr)
        {
            pressure_form = std::make_shared<PressureForm>(
                ndof,
                local_pressure_boundary,
                local_pressure_cavity,
                boundary_nodes,
                n_boundary_samples, *pressure_assembler,
                is_time_dependent);
            pressure_form->update_quantities(t, sol);
            forms.push_back(pressure_form);
        }
 
        inertia_form = nullptr;
        damping_form = nullptr;
        if (is_time_dependent)
        {
            if (!ignore_inertia)
            {
                assert(time_integrator != nullptr);
                inertia_form = std::make_shared<InertiaForm>(mass, *time_integrator);
                forms.push_back(inertia_form);
            }
 
            if (damping_assembler != nullptr)
            {
                damping_form = std::make_shared<ElasticForm>(
                    n_bases, bases, geom_bases, *damping_assembler, ass_vals_cache, t, dt, is_volume,
                    0., ElementInversionCheck::Discrete, conservative_max_iter);
                forms.push_back(damping_form);
            }
        }
        else
        {
            if (lagged_regularization_weight > 0)
            {
                forms.push_back(std::make_shared<LaggedRegForm>(lagged_regularization_iterations));
                forms.back()->set_weight(lagged_regularization_weight);
            }
        }
 
        if (rhs_assembler != nullptr)
        {
            // assembler::Mass mass_mat_assembler;
            // mass_mat_assembler.set_size(dim);
            StiffnessMatrix mass_tmp = mass;
            // mass_mat_assembler.assemble(dim == 3, n_bases, bases, geom_bases, mass_ass_vals_cache, mass_tmp, true);
            // assert(mass_tmp.rows() == mass.rows() && mass_tmp.cols() == mass.cols());
 
            if (!boundary_nodes.empty())
                al_form.push_back(std::make_shared<BCLagrangianForm>(
                    ndof, boundary_nodes, local_boundary, local_neumann_boundary,
                    n_boundary_samples, mass_tmp, *rhs_assembler, obstacle_ndof, is_time_dependent, t,
                    al_lumping));
            // forms.push_back(al_form.back());
        }
 
        if (!periodic_conditions.empty())
        {
            if (periodic_mesh == nullptr || periodic_local_boundary == nullptr)
                log_and_throw_error("Periodic boundary constraints require mesh boundary data");
 
            for (const json &condition : periodic_conditions)
            {
                const int condition_fe_space = condition.value("fe_space", -1);
                // -1 means all/unspecified FE spaces. If both IDs are explicit,
                // instantiate the condition only for the matching space.
                if (condition_fe_space >= 0 && fe_space_id >= 0 && condition_fe_space != fe_space_id)
                    continue;
 
                if (!condition.contains("boundary_ids") || !condition["boundary_ids"].is_array() || condition["boundary_ids"].size() != 2)
                    log_and_throw_error("A periodic boundary condition must contain exactly two boundary_ids");
 
                const std::array<int, 2> boundary_ids = {{condition["boundary_ids"][0].get<int>(),
                                                          condition["boundary_ids"][1].get<int>()}};
                al_form.push_back(std::make_shared<PeriodicBoundaryLagrangianForm>(
                    ndof, dim, *periodic_mesh, bases, *periodic_local_boundary,
                    boundary_ids, condition.value("tolerance", 1e-5)));
            }
        }
 
        bool add_zero_mean = false;
        if (zero_mean.is_boolean())
        {
            add_zero_mean = zero_mean.get<bool>();
        }
        else if (zero_mean.is_array())
        {
            const int current_fe_space = fe_space_id < 0 ? 0 : fe_space_id;
            add_zero_mean = std::find(zero_mean.begin(), zero_mean.end(), current_fe_space) != zero_mean.end();
        }
        else
        {
            log_and_throw_error("constraints.zero_mean must be a boolean or a list of FE-space IDs");
        }
 
        if (add_zero_mean)
        {
            StiffnessMatrix zero_mean_mass = mass;
            if (zero_mean_mass.rows() != ndof || zero_mean_mass.cols() != ndof)
            {
                assembler::Mass mass_assembler;
                mass_assembler.set_size(dim);
                mass_assembler.assemble(
                    dim == 3, n_bases, bases, geom_bases,
                    mass_ass_vals_cache, 0, zero_mean_mass, true);
            }
            if (zero_mean_mass.rows() != ndof || zero_mean_mass.cols() != ndof)
                log_and_throw_error("Unable to assemble zero-mean constraints for {} DoFs", ndof);
 
            const Eigen::VectorXd weights = zero_mean_mass * Eigen::VectorXd::Ones(ndof);
            std::vector<Eigen::Triplet<double>> entries;
            for (int d = 0; d < dim; ++d)
            {
                double weight_sum = 0;
                for (int dof = d; dof < ndof; dof += dim)
                    weight_sum += std::abs(weights(dof));
                if (weight_sum <= 0)
                    log_and_throw_error("Unable to assemble a zero-mean constraint for component {}", d);
                for (int dof = d; dof < ndof; dof += dim)
                    if (weights(dof) != 0)
                        entries.emplace_back(d, dof, weights(dof) / weight_sum);
            }
 
            StiffnessMatrix A(dim, ndof);
            A.setFromTriplets(entries.begin(), entries.end());
            Eigen::MatrixXd b = Eigen::MatrixXd::Zero(dim, 1);
            al_form.push_back(std::make_shared<MatrixLagrangianForm>(A, b));
        }
 
        for (const auto &path : hard_constraint_files)
        {
            logger().debug("Setting up hard constraints for {}", path);
            with_constraint_resource(resources, path, [&](const io::ResourceIO &file, const std::string &group) {
                const auto child = [&](const std::string &name) { return constraint_child(group, name); };
                std::vector<int> local2global;
                if (file.exists(child("local2global")))
                    local2global = file.read_int_vector(child("local2global"));
 
                if (local2global.empty())
                {
                    local2global.resize(in_node_to_node.size());
                    for (int i = 0; i < local2global.size(); ++i)
                        local2global[i] = in_node_to_node[i];
                }
                else
                {
                    for (auto &v : local2global)
                        v = in_node_to_node[v];
                }
 
                const Eigen::MatrixXd bin = file.read_matrix(child("b"));
                StiffnessMatrix A, A_proj;
                Eigen::MatrixXd b, b_proj;
 
                if (file.exists(child("A")))
                {
                    const Eigen::MatrixXd Ain = file.read_matrix(child("A"));
                    utils::scatter_matrix(ndof, dim, Ain, bin, local2global, A, b);
                    if (file.exists(child("A_proj")))
                    {
                        if (!file.exists(child("b_proj")))
                            log_and_throw_error("Missing b_proj in hard constraint resource");
                        utils::scatter_matrix_col(
                            ndof, dim, file.read_matrix(child("A_proj")), file.read_matrix(child("b_proj")),
                            local2global, A_proj, b_proj);
                    }
                }
                else
                {
                    utils::scatter_matrix(
                        ndof, dim, file.read_long_vector(child("A_triplets/shape")),
                        file.read_int_vector(child("A_triplets/rows")), file.read_int_vector(child("A_triplets/cols")),
                        file.read_double_vector(child("A_triplets/values")), bin, local2global, A, b);
 
                    if (file.is_group(child("A_proj_triplets")))
                    {
                        for (const std::string &required : {"b_proj", "A_proj_triplets/rows", "A_proj_triplets/cols", "A_proj_triplets/values", "A_proj_triplets/shape"})
                            if (!file.exists(child(required)))
                                log_and_throw_error("Missing {} in hard constraint resource", required);
                        utils::scatter_matrix_col(
                            ndof, dim, file.read_long_vector(child("A_proj_triplets/shape")),
                            file.read_int_vector(child("A_proj_triplets/rows")), file.read_int_vector(child("A_proj_triplets/cols")),
                            file.read_double_vector(child("A_proj_triplets/values")), file.read_matrix(child("b_proj")),
                            local2global, A_proj, b_proj);
                    }
                }
 
                al_form.push_back(std::make_shared<MatrixLagrangianForm>(A, b, A_proj, b_proj));
            });
            // forms.push_back(al_form.back());
        }
 
        for (const auto &j : soft_constraint_files)
        {
            const std::string &path = j["data"];
            double weight = j["weight"];
 
            logger().debug("Setting up soft constraints for {}", path);
            with_constraint_resource(resources, path, [&](const io::ResourceIO &file, const std::string &group) {
                const auto child = [&](const std::string &name) { return constraint_child(group, name); };
                std::vector<int> local2global;
                if (file.exists(child("local2global")))
                    local2global = file.read_int_vector(child("local2global"));
                if (local2global.empty())
                {
                    local2global.resize(in_node_to_node.size());
                    for (int i = 0; i < local2global.size(); ++i)
                        local2global[i] = in_node_to_node[i];
                }
                else
                {
                    for (auto &v : local2global)
                        v = in_node_to_node[v];
                }
 
                const Eigen::MatrixXd bin = file.read_matrix(child("b"));
                StiffnessMatrix A;
                Eigen::MatrixXd b;
                if (file.exists(child("A")))
                    utils::scatter_matrix(ndof, dim, file.read_matrix(child("A")), bin, local2global, A, b);
                else
                    utils::scatter_matrix(
                        ndof, dim, file.read_long_vector(child("A_triplets/shape")),
                        file.read_int_vector(child("A_triplets/rows")), file.read_int_vector(child("A_triplets/cols")),
                        file.read_double_vector(child("A_triplets/values")), bin, local2global, A, b);
 
                forms.push_back(std::make_shared<QuadraticPenaltyForm>(A, b, weight));
            });
        }
 
        if (macro_strain_constraint.is_active())
        {
            // don't push these two into forms because they take a different input x
            strain_al_lagr_form = std::make_shared<MacroStrainLagrangianForm>(macro_strain_constraint);
        }
 
        contact_form = nullptr;
        periodic_contact_form = nullptr;
        friction_form = nullptr;
        if (contact_enabled)
        {
            const bool use_adaptive_barrier_stiffness = !barrier_stiffness.is_number();
 
            if (periodic_contact)
            {
                periodic_contact_form = std::make_shared<PeriodicContactForm>(
                    collision_mesh, tiled_to_single, dhat, avg_mass, use_area_weighting, use_improved_max_operator, use_physical_barrier,
                    use_adaptive_barrier_stiffness, is_time_dependent, enable_shape_derivatives, broad_phase, ccd_tolerance,
                    ccd_max_iterations);
 
                if (use_adaptive_barrier_stiffness)
                {
                    periodic_contact_form->set_barrier_stiffness(1);
                    // logger().debug("Using adaptive barrier stiffness");
                }
                else
                {
                    assert(barrier_stiffness.is_number());
                    assert(barrier_stiffness.get<double>() > 0);
                    periodic_contact_form->set_barrier_stiffness(barrier_stiffness);
                    // logger().debug("Using fixed barrier stiffness of {}", contact_form->barrier_stiffness());
                }
 
                // periodic_contact_form is not pushed into forms since it takes different input vectors.
            }
            else
            {
                if (use_gcp_formulation)
                {
                    contact_form = std::make_shared<SmoothContactForm>(
                        collision_mesh, dhat, avg_mass, alpha_t, alpha_n, use_adaptive_dhat, min_distance_ratio,
                        use_adaptive_barrier_stiffness, is_time_dependent, enable_shape_derivatives, broad_phase,
                        ccd_tolerance * units.characteristic_length(), ccd_max_iterations);
                }
                else
                {
                    contact_form = std::make_shared<BarrierContactForm>(
                        collision_mesh, dhat, avg_mass, use_area_weighting, use_improved_max_operator, use_physical_barrier,
                        use_adaptive_barrier_stiffness, is_time_dependent, enable_shape_derivatives, broad_phase, ccd_tolerance * units.characteristic_length(),
                        ccd_max_iterations);
                }
 
                if (use_adaptive_barrier_stiffness)
                {
                    contact_form->set_barrier_stiffness(initial_barrier_stiffness);
                    // logger().debug("Using adaptive barrier stiffness");
                }
                else
                {
                    assert(barrier_stiffness.is_number());
                    assert(barrier_stiffness.get<double>() > 0);
                    contact_form->set_barrier_stiffness(barrier_stiffness);
                    // logger().debug("Using fixed barrier stiffness of {}", contact_form->barrier_stiffness());
                }
 
                forms.push_back(contact_form);
            }
 
            if (friction_coefficient != 0)
            {
                friction_form = std::make_shared<FrictionForm>(
                    collision_mesh, time_integrator, epsv, friction_coefficient,
                    broad_phase, *contact_form, friction_iterations);
                friction_form->init_lagging(sol);
                forms.push_back(friction_form);
            }
 
            if (adhesion_enabled)
            {
                normal_adhesion_form = std::make_shared<NormalAdhesionForm>(
                    collision_mesh, dhat_p, dhat_a, Y, is_time_dependent, enable_shape_derivatives,
                    broad_phase, ccd_tolerance * units.characteristic_length(), ccd_max_iterations);
                forms.push_back(normal_adhesion_form);
 
                if (tangential_adhesion_coefficient != 0)
                {
                    tangential_adhesion_form = std::make_shared<TangentialAdhesionForm>(
                        collision_mesh, time_integrator, epsa, tangential_adhesion_coefficient,
                        broad_phase, *normal_adhesion_form, tangential_adhesion_iterations);
                    forms.push_back(tangential_adhesion_form);
                }
            }
        }
 
        const std::vector<json> rayleigh_damping_jsons = utils::json_as_array(rayleigh_damping);
        if (is_time_dependent)
        {
            // Map from form name to form so RayleighDampingForm::create can get the correct form to damp
            const std::unordered_map<std::string, std::shared_ptr<Form>> possible_forms_to_damp = {
                {"elasticity", elastic_form},
                {"contact", contact_form},
            };
 
            for (const json &params : rayleigh_damping_jsons)
            {
                forms.push_back(RayleighDampingForm::create(
                    params, possible_forms_to_damp,
                    *time_integrator));
            }
        }
        else if (rayleigh_damping_jsons.size() > 0)
        {
            log_and_throw_adjoint_error("Rayleigh damping is only supported for time-dependent problems");
        }
 
        update_dt();
 
        return forms;
    }
 
    void SolveData::update_barrier_stiffness(const Eigen::VectorXd &x)
    {
        if (contact_form == nullptr || !contact_form->use_adaptive_barrier_stiffness())
            return;
 
        Eigen::VectorXd grad_energy = Eigen::VectorXd::Zero(x.size());
        const std::array<std::shared_ptr<Form>, 4> energy_forms{
            {elastic_form, inertia_form, body_form, pressure_form}};
        for (const std::shared_ptr<Form> &form : energy_forms)
        {
            if (form == nullptr || !form->enabled())
                continue;
 
            Eigen::VectorXd grad_form;
            form->first_derivative(x, grad_form);
            grad_energy += grad_form;
        }
 
        contact_form->update_barrier_stiffness(x, grad_energy);
    }
 
    void SolveData::update_dt()
    {
        if (time_integrator == nullptr) // if is not time dependent
            return;
 
        const std::array<std::shared_ptr<Form>, 6> energy_forms{
            {elastic_form, body_form, pressure_form, damping_form, contact_form, friction_form}};
        for (const std::shared_ptr<Form> &form : energy_forms)
        {
            if (form == nullptr)
                continue;
            form->set_weight(time_integrator->acceleration_scaling());
        }
    }
 
    std::vector<std::pair<std::string, std::shared_ptr<solver::Form>>> SolveData::named_forms() const
    {
        std::vector<std::pair<std::string, std::shared_ptr<solver::Form>>> res{
            {"elastic", elastic_form},
            {"inertia", inertia_form},
            {"body", body_form},
            {"contact", contact_form},
            {"friction", friction_form},
            {"damping", damping_form},
            {"pressure", pressure_form},
            {"strain_augmented_lagrangian_lagr", strain_al_lagr_form},
            {"periodic_contact", periodic_contact_form},
        };
 
        for (const auto &form : al_form)
            res.push_back({"augmented_lagrangian", form});
 
        return res;
    }
} // namespace polyfem::solver