diff --git a/free_energy.cpp b/free_energy.cpp new file mode 100644 index 0000000..e69de29 diff --git a/mapping.cpp b/mapping.cpp new file mode 100644 index 0000000..84da56d --- /dev/null +++ b/mapping.cpp @@ -0,0 +1,2103 @@ +//region Includes +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +//endregion + +//region Test Utilities +constexpr std::string_view ANSI_GREEN = "\033[32m"; +constexpr std::string_view ANSI_RED = "\033[31m"; +constexpr std::string_view ANSI_YELLOW = "\033[33m"; +constexpr std::string_view ANSI_BLUE = "\033[34m"; +constexpr std::string_view ANSI_RESET = "\033[0m"; + +enum class TEST_RESULT_TYPE : uint8_t { + SUCCESS, + FAILURE, + PARTIAL +}; + +std::string fmt_test_msg(const std::string_view test_name, const TEST_RESULT_TYPE type, size_t num_fails, size_t total) { + std::string_view color; + switch (type) { + case TEST_RESULT_TYPE::SUCCESS: + color = ANSI_GREEN; + break; + case TEST_RESULT_TYPE::FAILURE: + color = ANSI_RED; + break; + case TEST_RESULT_TYPE::PARTIAL: + color = ANSI_YELLOW; + break; + default: + color = ANSI_RESET; + } + return std::format("{}[TEST: {}] {}/{}{}", color, test_name, total-num_fails, total, ANSI_RESET); +} +//endregion + +//region Constants +/******************** + * Constants + *********************/ +constexpr double G = 1.0; +constexpr double MASS = 1.0; +constexpr double RADIUS = 1.0; +constexpr double CENTRAL_DENSITY = 1.0; + +constexpr char HOST[10] = "localhost"; +constexpr int PORT = 19916; + +//endregion + +//region Concepts and Typedefs +/******************** + * Concepts + *********************/ +template +concept is_xad = + std::is_same_v> + || std::is_same_v> + || std::is_same_v>; + +template +concept is_real = std::is_floating_point_v || is_xad; + +/******************** + * Type Defs + *********************/ +template +using EOS_P = std::function; +//endregion + +//region User Argument Structs +/******************** + * User Args + *********************/ +struct potential { + double tol; + int max_iters; +}; + +struct rot { + bool enabled; + double omega; +}; + +struct Args { + std::string mesh_file; + potential p{}; + rot r{}; + bool verbose{}; + double index{}; + double mass{}; + double c{}; + + int max_iters{}; + double tol{}; +}; +//endregion + +//region Misc Structs +struct OblatePotential { + bool use{false}; + double a{1}; + double c{1}; + double rho_0{1}; +}; + +struct Bounds { + double r_star_ref; + double r_inf_ref; +}; + +enum BoundsError : uint8_t { + CANNOT_FIND_VACUUM +}; +//endregion + +//region Domain Enums +enum class Domains : uint8_t { + NONE = 0, + CORE = 1 << 0, + ENVELOPE = 1 << 1, + VACUUM = 1 << 2, + ALL = 0x7 +}; + +inline Domains operator|(Domains lhs, Domains rhs) { + return static_cast(static_cast(lhs) | static_cast(rhs)); +} + +inline Domains operator&(Domains lhs, Domains rhs) { + return static_cast(static_cast(lhs) & static_cast(rhs)); +} + +const Domains STELLAR = Domains::CORE | Domains::ENVELOPE; +//endregion + +//region Domain Mapper +/******************** + * Mappers + *********************/ +class DomainMapper { +public: + DomainMapper( + const double r_star_ref, + const double r_inf_ref + ) : + m_d(nullptr), + m_r_star_ref(r_star_ref), + m_r_inf_ref(r_inf_ref) {} + + explicit DomainMapper( + const mfem::GridFunction &d, + const double r_star_ref, + const double r_inf_ref + ) : + m_d(&d), + m_dim(d.FESpace()->GetMesh()->Dimension()), + m_r_star_ref(r_star_ref), + m_r_inf_ref(r_inf_ref) {}; + + [[nodiscard]] bool is_vacuum(const mfem::ElementTransformation &T) const { + if (T.ElementType == mfem::ElementTransformation::ELEMENT) { + return T.Attribute == m_vacuum_attr; + } else if (T.ElementType == mfem::ElementTransformation::BDR_ELEMENT) { + return T.Attribute == m_vacuum_attr - 1; // TODO: In a more robust code this should really be read from the stroid API to ensure that the vacuum boundary is really 1 - the vacuum material attribute + } + return false; + } + + void SetDisplacement(const mfem::GridFunction &d) { + if (m_dim != d.FESpace()->GetMesh()->Dimension()) { + const std::string err_msg = std::format("Dimension mismatch: DomainMapper is initialized for dimension {}, but provided displacement field has dimension {}.", m_dim, d.FESpace()->GetMesh()->Dimension()); + throw std::invalid_argument(err_msg); + } + m_d = &d; + } + + [[nodiscard]] bool IsIdentity() const { + return (m_d == nullptr); + } + + void ResetDisplacement() { + m_d = nullptr; + } + + void ComputeJacobian(mfem::ElementTransformation &T, mfem::DenseMatrix &J) const { + mfem::DenseMatrix J_D(m_dim, m_dim); + J.SetSize(m_dim, m_dim); + J = 0.0; + J_D = 0.0; + if (IsIdentity()) { + for (int i = 0; i < m_dim; ++i) { + J_D(i, i) = 1.0; // Identity mapping + } + } else { + m_d->GetVectorGradient(T, J_D); + for (int i = 0; i < m_dim; ++i) { + J_D(i, i) += 1.0; // Add identity to get the total Jacobian of the mapping + } + } + + if (is_vacuum(T)) { + mfem::Vector x_ref(m_dim); + T.Transform(T.GetIntPoint(), x_ref); + + mfem::Vector d_val(m_dim), x_disp(m_dim); + if (IsIdentity()) { + x_disp = x_ref; + } else { + m_d->GetVectorValue(T, T.GetIntPoint(), d_val); + add(x_ref, d_val, x_disp); + } + + mfem::DenseMatrix J_K(m_dim, m_dim); + ComputeKelvinJacobian(x_ref, x_disp, J_K); + mfem::Mult(J_K, J_D, J); + } else { + J = J_D; + } + } + + double ComputeDetJ(mfem::ElementTransformation& T, const mfem::IntegrationPoint& ip) const { + if (IsIdentity() && !is_vacuum(T)) return 1.0; // If no mapping, the determinant of the Jacobian is 1 + mfem::DenseMatrix J; + ComputeJacobian(T, J); + return J.Det(); + } + + void ComputeMappedDiffusionTensor(mfem::ElementTransformation &T, mfem::DenseMatrix &D) const { + mfem::DenseMatrix J(m_dim, m_dim), JInv(m_dim, m_dim); + + ComputeJacobian(T, J); + const double detJ = J.Det(); + mfem::CalcInverse(J, JInv); + D.SetSize(m_dim, m_dim); + mfem::MultABt(JInv, JInv, D); + D *= fabs(detJ); + } + + void ComputeInverseJacobian(mfem::ElementTransformation &T, mfem::DenseMatrix &JInv) const { + mfem::DenseMatrix J(m_dim, m_dim); + ComputeJacobian(T, J); + JInv.SetSize(m_dim, m_dim); + mfem::CalcInverse(J, JInv); + } + + void GetPhysicalPoint(mfem::ElementTransformation& T, const mfem::IntegrationPoint& ip, mfem::Vector& x_phys) const { + x_phys.SetSize(m_dim); + mfem::Vector x_ref(m_dim); + T.Transform(ip, x_ref); + + if (IsIdentity()) { + x_phys = x_ref; + } else { + mfem::Vector d_val(m_dim); + m_d->GetVectorValue(T, ip, d_val); + add(x_ref, d_val, x_phys); + } + if (is_vacuum(T)) { + ApplyKelvinMapping(x_ref, x_phys); + } + } + + [[nodiscard]] const mfem::GridFunction* GetDisplacement() const { return m_d; } + + +private: + void ApplyKelvinMapping(const mfem::Vector& x_ref, mfem::Vector& x_phys) const { + const double r_ref = x_ref.Norml2(); + double xi = (r_ref - m_r_star_ref) / (m_r_inf_ref - m_r_star_ref); + xi = std::clamp(xi, 0.0, m_xi_clamp); + const double factor = m_r_star_ref / (r_ref * (1 - xi)); + x_phys *= factor; + } + + void ComputeKelvinJacobian(const mfem::Vector& x_ref, const mfem::Vector &x_disp, mfem::DenseMatrix& J_K) const { + const double r_ref = x_ref.Norml2(); + const double delta_R = m_r_inf_ref - m_r_star_ref; + + double xi = (r_ref - m_r_star_ref) / delta_R; + xi = std::clamp(xi, 0.0, m_xi_clamp); + + const double denom = 1.0 - xi; + + const double k = m_r_star_ref / (r_ref * denom); + + const double dk_dr = m_r_star_ref * (( 1.0 / (delta_R* r_ref * denom * denom)) - ( 1.0 / (r_ref * r_ref * denom))); + + J_K.SetSize(m_dim, m_dim); + const double outer_factor = dk_dr / r_ref; + + for (int i = 0; i < m_dim; ++i) { + for (int j = 0; j < m_dim; ++j) { + J_K(i, j) = outer_factor * x_disp(i) * x_ref(j); + if (i == j) { + J_K(i, j) += k; + } + } + } + } +private: + const mfem::GridFunction *m_d; + std::unique_ptr m_internal_d; + const int m_dim{3}; + const int m_vacuum_attr{3}; + const double m_r_star_ref{1.0}; + const double m_r_inf_ref{2.0}; + const double m_xi_clamp{0.999}; +}; + +//endregion + +//region State Types +/******************** + * State Types + *********************/ +struct FEM { + std::unique_ptr mesh; + std::unique_ptr H1_fec; + std::unique_ptr H1_fes; + std::unique_ptr Vec_H1_fes; + std::unique_ptr mapping; + + mfem::Array block_offsets; + std::unique_ptr reference_x; + + mfem::Vector com; + mfem::DenseMatrix Q; + mfem::Array ess_tdof_x; + + std::unique_ptr> star_marker; + std::unique_ptr> vacuum_marker; + std::unique_ptr> surface_marker; + std::unique_ptr> inf_marker; + + [[nodiscard]] bool okay() const { return (mesh != nullptr) && (H1_fec != nullptr) && (H1_fes != nullptr) && (Vec_H1_fes != nullptr); } + + [[nodiscard]] bool has_mapping() const { return mapping != nullptr; } +}; + +struct CoupledState { + std::unique_ptr U; + mfem::GridFunction rho; + mfem::GridFunction d; // Stability depends on solving for the displacement vector not the nodal position vector, those live on a reference domain. + + explicit CoupledState(const FEM& fem) { + U = std::make_unique(fem.block_offsets); + rho.MakeRef(fem.H1_fes.get(), U->GetBlock(0), 0); + d.MakeRef(fem.Vec_H1_fes.get(), U->GetBlock(1), 0); + *U = 0.0; + U->GetBlock(2) = 1.0; + } +}; + +//endregion + +//region Function Definitions +/******************** + * Core Setup Functions + *********************/ +FEM setup_fem(const std::string& filename, bool verbose=true); + +/******************** + * Utility Functions + *********************/ +void view_mesh(const std::string& host, int port, const mfem::Mesh& mesh, const mfem::GridFunction& gf, const std::string& title); +double domain_integrate_grid_function(const FEM& fem, const mfem::GridFunction& gf, Domains domain = Domains::ALL); +mfem::Vector get_com(const FEM& fem, const mfem::GridFunction &rho); +void get_physical_coordinates(const mfem::GridFunction& reference_pos, const mfem::GridFunction& displacement, mfem::GridFunction& physical_pos); +void populate_element_mask(const FEM& fem, Domains domain, mfem::Array& mask); +std::expected DiscoverBounds(const mfem::Mesh *mesh, int vacuum_attr); + +/******************** + * Physics Functions + *********************/ +double centrifugal_potential(const mfem::Vector& phys_x, double omega); +double get_moment_of_inertia(const FEM& fem, const mfem::GridFunction& rho); +double oblate_spheroid_surface_potential(const mfem::Vector& x, double a, double c, double total_mass); +std::unique_ptr grav_potential(const FEM& fem, const Args &args, const mfem::GridFunction& rho); +std::unique_ptr get_potential(const FEM& fem, const Args& args, const mfem::GridFunction& rho); +mfem::DenseMatrix compute_quadrupole_moment_tensor(const FEM& fem, const mfem::GridFunction& rho, const mfem::Vector& com); +double l2_multipole_potential(const FEM& fem, double total_mass, const mfem::Vector& phys_x); + +/******************** + * Conservation Functions + *********************/ +void conserve_mass(const FEM& fem, mfem::GridFunction& rho, double target_mass); +//endregion + +//region Mapping Coefficients +class MappedScalarCoefficient : public mfem::Coefficient { +public: + enum class EVAL_POINTS : uint8_t { + PHYSICAL, + REFERENCE + }; + + MappedScalarCoefficient( + const DomainMapper& map, + mfem::Coefficient& coeff, + const EVAL_POINTS eval_point=EVAL_POINTS::PHYSICAL + ) : + m_map(map), + m_coeff(coeff), + m_eval_point(eval_point) {}; + + double Eval(mfem::ElementTransformation &T, const mfem::IntegrationPoint &ip) override { + T.SetIntPoint(&ip); + + double detJ = m_map.ComputeDetJ(T, ip); + double f_val = 0.0; + + switch (m_eval_point) { + case EVAL_POINTS::PHYSICAL: { + f_val = eval_at_point(m_coeff, T, ip); + break; + } + case EVAL_POINTS::REFERENCE: { + f_val = m_coeff.Eval(T, ip); + break; + } + } + return f_val * fabs(detJ); + } +private: + static double eval_at_point(mfem::Coefficient& c, mfem::ElementTransformation& T, const mfem::IntegrationPoint& ip) { + return c.Eval(T, ip); + } + +private: + const DomainMapper& m_map; + mfem::Coefficient& m_coeff; + EVAL_POINTS m_eval_point; + +}; + +class MappedDiffusionCoefficient : public mfem::MatrixCoefficient { +public: + MappedDiffusionCoefficient( + const DomainMapper& map, + mfem::Coefficient& sigma, + const int dim + ) : + mfem::MatrixCoefficient(dim), + m_map(map), + m_scalar(&sigma), + m_tensor(nullptr) {}; + + MappedDiffusionCoefficient( + const DomainMapper& map, + mfem::MatrixCoefficient& sigma + ) : + mfem::MatrixCoefficient(sigma.GetHeight()), + m_map(map), + m_scalar(nullptr), + m_tensor(&sigma) {}; + + void Eval(mfem::DenseMatrix &K, mfem::ElementTransformation &T, const mfem::IntegrationPoint &ip) override { + const int dim = height; + T.SetIntPoint(&ip); + + mfem::DenseMatrix J(dim, dim), JInv(dim, dim); + m_map.ComputeJacobian(T, J); + const double detJ = J.Det(); + mfem::CalcInverse(J, JInv); + + if (m_scalar) { + const double sig_val = m_scalar->Eval(T, ip); + mfem::MultABt(JInv, JInv, K); + K *= sig_val * fabs(detJ); + } else { + mfem::DenseMatrix sig_mat(dim, dim); + m_tensor->Eval(sig_mat, T, ip); + + mfem::DenseMatrix temp(dim, dim); + Mult(JInv, sig_mat, temp); + + MultABt(temp, JInv, K); + K *= fabs(detJ); + } + } +private: + const DomainMapper& m_map; + mfem::Coefficient* m_scalar; + mfem::MatrixCoefficient* m_tensor; +}; + +class MappedVectorCoefficient : public mfem::VectorCoefficient { +public: + MappedVectorCoefficient( + const DomainMapper& map, + mfem::VectorCoefficient& coeff + ) : + mfem::VectorCoefficient(coeff.GetVDim()), + m_map(map), + m_coeff(coeff) {}; + + void Eval(mfem::Vector& V, mfem::ElementTransformation& T, const mfem::IntegrationPoint& ip) override { + const int dim = vdim; + T.SetIntPoint(&ip); + + mfem::DenseMatrix JInv(dim, dim); + m_map.ComputeInverseJacobian(T, JInv); + double detJ = m_map.ComputeDetJ(T, ip); + + mfem::Vector C_phys(dim); + m_coeff.Eval(C_phys, T, ip); + + V.SetSize(dim); + JInv.Mult(C_phys, V); + V *= fabs(detJ); + } +private: + const DomainMapper& m_map; + mfem::VectorCoefficient& m_coeff; +}; + +class PhysicalPositionFunctionCoefficient : public mfem::Coefficient { +public: + using Func = std::function; + + PhysicalPositionFunctionCoefficient( + const DomainMapper& map, + Func f + ) : + m_f(std::move(f)), + m_map(map) {}; + + double Eval(mfem::ElementTransformation &T, const mfem::IntegrationPoint &ip) override { + T.SetIntPoint(&ip); + mfem::Vector x; + m_map.GetPhysicalPoint(T, ip, x); + return m_f(x); + } +private: + Func m_f; + const DomainMapper& m_map; +}; +//endregion + +//region Integrators +/******************** + * Integrators + *********************/ +template +class FluidIntegrator : public mfem::NonlinearFormIntegrator { + using Scalar = EOS_T::value_type; +public: + explicit FluidIntegrator( + EOS_P eos, + const DomainMapper* map = nullptr + ) : + m_eos(std::move(eos)), + m_map(map) + {}; + + void AssembleElementVector( + const mfem::FiniteElement &el, + mfem::ElementTransformation &Tr, + const mfem::Vector &elfun, + mfem::Vector &elvect + ) override { + const int dof = el.GetDof(); + elvect.SetSize(dof); + elvect = 0.0; + + const mfem::IntegrationRule *ir = IntRule ? IntRule : &mfem::IntRules.Get(el.GetGeomType(), 2 * el.GetOrder() + 1); + + mfem::Vector shape(dof); + for (int i = 0; i < ir->GetNPoints(); i++) { + const mfem::IntegrationPoint& ip = ir->IntPoint(i); + Tr.SetIntPoint(&ip); + el.CalcShape(ip, shape); + + double u = shape * elfun; + EOS_T rho = u; + const double val = m_eos(rho, 0.0).value(); + double weight = ip.weight * Tr.Weight() * val; + + if (m_map) weight *= fabs(m_map->ComputeDetJ(Tr, ip)); + + + elvect.Add(weight, shape); + } + } + void AssembleElementGrad( + const mfem::FiniteElement &el, + mfem::ElementTransformation &Tr, + const mfem::Vector &elfun, + mfem::DenseMatrix &elmat + ) override { + const int dof = el.GetDof(); + elmat.SetSize(dof); + elmat = 0.0; + + const mfem::IntegrationRule *ir = IntRule ? IntRule : &mfem::IntRules.Get(el.GetGeomType(), 2 * el.GetOrder() + 1); + + mfem::Vector shape(dof); + for (int i = 0; i < ir->GetNPoints(); i++) { + const mfem::IntegrationPoint& ip = ir->IntPoint(i); + Tr.SetIntPoint(&ip); + el.CalcShape(ip, shape); + + double u = shape * elfun; + + double d_val_d_rho = 0.0; + if (u > 1e-15) { + xad::Tape tape; + EOS_T x_r = u; + EOS_T x_t = 0.0; // In future this is one area where we introduce a temp dependency + + tape.registerInput(x_r); + EOS_T result = m_eos(x_r, x_t); + tape.registerOutput(result); + result.setAdjoint(1.0); + tape.computeAdjoints(); + d_val_d_rho = x_r.getAdjoint(); + } + + double weight = ip.weight * Tr.Weight() * d_val_d_rho; + if (m_map) weight *= fabs(m_map->ComputeDetJ(Tr, ip)); + + mfem::AddMult_a_VVt(weight, shape, elmat); + } + } + + [[nodiscard]] bool has_mapping() const { return m_map != nullptr; } + + void set_mapping(const DomainMapper* map) { m_map = map; } + + void clear_mapping() { m_map = nullptr; } +private: + EOS_P m_eos; + const DomainMapper* m_map{nullptr}; +}; + +//endregion + +//region Coefficients +/******************** + * Coefficient Defs + *********************/ +template +class PressureBoundaryForce : public mfem::VectorCoefficient { +public: + PressureBoundaryForce( + const int dim, + const FEM& fem, + const mfem::GridFunction& rho, + const EOS_P& eos, + const double P_fit + ) : VectorCoefficient(dim), m_fem(fem), m_rho(rho), m_eos(eos), m_P_fit(P_fit) {}; + + void Eval(mfem::Vector &V, mfem::ElementTransformation &T, const mfem::IntegrationPoint &ip) override { + V.SetSize(vdim); + V = 0.0; + + double rho = m_rho.GetValue(T, ip); + + EOS_T x_rho = rho; + const double P_curr = m_eos(x_rho, 0.0).value(); + const double delta_P = P_curr - m_P_fit; + + mfem::Vector phys(vdim); + T.Transform(ip, phys); + mfem::Vector normal(vdim); + mfem::CalcOrtho(T.Jacobian(), normal); + + for (int i = 0; i < vdim; ++i) { + V(i) = delta_P * normal(i); + } + } +private: + const FEM& m_fem; + const mfem::GridFunction& m_rho; + const EOS_P& m_eos; + double m_P_fit; +}; + +template +class MappingDetJacobianCoefficient : public mfem::Coefficient { +public: + MappingDetJacobianCoefficient( + const mfem::GridFunction& d, + const mfem::FiniteElementSpace& fes + ) : + m_d(d), + m_fes(fes) {} + + double Eval(mfem::ElementTransformation &T, const mfem::IntegrationPoint &ip) override { + const int dim = T.GetSpaceDim(); + mfem::DenseMatrix grad_d(dim, dim); + m_d.GetVectorGradient(T, grad_d); + for (int i = 0; i < dim; ++i) { + grad_d(i, i) += 1.0; + } + return grad_d.Det(); + } +private: + const mfem::GridFunction& m_d; // Displacement vector + const mfem::FiniteElementSpace &m_fes; +}; +//endregion + +//region Operators +/******************** + * Operator Defs + *********************/ +template +class VectorOperator : public mfem::Operator { +public: + VectorOperator() = default; + + VectorOperator( + const mfem::Vector& v, + const bool is_col + ) : + Operator(is_col ? v.Size() : 1, is_col ? 1 : v.Size()), + m_v(v), + m_is_col(is_col), + m_is_initialized(true) {} + + void SetVector(const mfem::Vector& v, const bool is_col) { + if (v.Size() != m_v.Size()) { + m_v.SetSize(v.Size()); + } + m_v = v; + m_is_col = is_col; + height = is_col ? v.Size() : 1; + width = is_col ? 1 : v.Size(); + m_is_initialized = true; + } + + void Mult(const mfem::Vector &x, mfem::Vector &y) const override { + if (!m_is_initialized) throw std::runtime_error("VectorOperator Not initialized"); + if (m_is_col) { + y.SetSize(m_v.Size()); + y = 0.0; + y.Add(x(0), m_v); + } else { + y.SetSize(1); + y(0) = m_v * x; + } + } +private: + mfem::Vector m_v; + bool m_is_col{false}; + bool m_is_initialized{false}; +}; + +template +class PressureDensityCoupling : public mfem::Operator { +public: + PressureDensityCoupling( + FEM& fem, + const mfem::GridFunction& rho, + const EOS_P& eos_pressure) + : Operator(fem.Vec_H1_fes->GetTrueVSize(), fem.H1_fes->GetTrueVSize()), + m_fem(fem), + m_rho(rho), + m_eos(eos_pressure) { + Assemble(); + } + + void Assemble() { + const int dim = m_fem.mesh->Dimension(); + const int scalar_size = m_fem.H1_fes->GetTrueVSize(); + const int vector_size = m_fem.Vec_H1_fes->GetTrueVSize(); + + m_mat = std::make_unique(vector_size, scalar_size); + + for (int be = 0; be < m_fem.mesh->GetNBE(); ++be) { + auto* ftr = m_fem.mesh->GetBdrFaceTransformations(be); + if (!ftr) continue; + + const int elem = ftr->Elem1No; + + const auto& scalar_fe = *m_fem.H1_fes->GetFE(elem); + + const int sdof = scalar_fe.GetDof(); + + mfem::Array scalar_dofs, vector_dofs; + m_fem.H1_fes->GetElementDofs(elem, scalar_dofs); + m_fem.Vec_H1_fes->GetElementDofs(elem, vector_dofs); + + mfem::DenseMatrix elmat(vector_dofs.Size(), scalar_dofs.Size()); + elmat = 0.0; + + const auto& face_ir = mfem::IntRules.Get(ftr->GetGeometryType(), 2 * scalar_fe.GetOrder() + 2); + + mfem::Vector shape(sdof); + mfem::Vector normal(dim); + + for (int q = 0; q < face_ir.GetNPoints(); ++q) { + const auto& face_ip = face_ir.IntPoint(q); + + ftr->SetAllIntPoints(&face_ip); + + const mfem::IntegrationPoint& vol_ip = ftr->GetElement1IntPoint(); + + scalar_fe.CalcShape(vol_ip, shape); + + mfem::CalcOrtho(ftr->Face->Jacobian(), normal); + + mfem::ElementTransformation& vol_trans = ftr->GetElement1Transformation(); + vol_trans.SetIntPoint(&vol_ip); + double rho_val = m_rho.GetValue(vol_trans, vol_ip); + + double dPdrho = 0.0; + if (rho_val > 1e-15) { + using Scalar = EOS_T::value_type; + xad::Tape tape; + + EOS_T x_rho = rho_val; + tape.registerInput(x_rho); + EOS_T P = m_eos(x_rho, EOS_T(0.0)); + tape.registerOutput(P); + + P.setAdjoint(1.0); + + tape.computeAdjoints(); + dPdrho = x_rho.getAdjoint(); + } + + const double w = face_ip.weight; + + for (int k = 0; k < sdof; ++k) { + for (int j = 0; j < sdof; ++j) { + double base = -dPdrho * shape(j) * shape(k) * w; + for (int c = 0; c < dim; ++c) { + elmat(k + c * sdof, j) += base * normal(c); + } + } + } + } + m_mat->AddSubMatrix(vector_dofs, scalar_dofs, elmat); + } + m_mat->Finalize(); + } + + void Mult(const mfem::Vector &x, mfem::Vector &y) const override { + m_mat->Mult(x, y); + } + + [[nodiscard]] mfem::SparseMatrix& SpMat() const { return *m_mat; } + +private: + FEM& m_fem; + const mfem::GridFunction& m_rho; + const EOS_P& m_eos; + std::unique_ptr m_mat; +}; + +template +class MassDisplacementCoupling : public mfem::Operator { +public: + MassDisplacementCoupling( + FEM& fem, + const mfem::GridFunction& rho, + const bool is_col + ) : + Operator( + is_col ? fem.Vec_H1_fes->GetTrueVSize() : 1, + is_col ? 1 : fem.Vec_H1_fes->GetTrueVSize() + ), + m_fem(fem), + m_rho(rho), + m_is_col(is_col){ + m_D.SetSize(m_fem.Vec_H1_fes->GetTrueVSize()); + Assemble(); + } + + void Assemble() const { + const int dim = m_fem.mesh->Dimension(); + m_D = 0.0; + for (int elemID = 0; elemID < m_fem.mesh->GetNE(); ++elemID) { + auto* trans = m_fem.mesh->GetElementTransformation(elemID); + const auto& fe = *m_fem.Vec_H1_fes->GetFE(elemID); + const int dof = fe.GetDof(); + + mfem::Array vdofs; + m_fem.Vec_H1_fes->GetElementDofs(elemID, vdofs); + + const auto& ir = mfem::IntRules.Get(trans->GetGeometryType(), 2 * fe.GetOrder() + 1); + + mfem::DenseMatrix dshape(dof, dim); + mfem::Vector elvec(dof * dim); + elvec = 0.0; + + for (int q = 0; q < ir.GetNPoints(); ++q) { + const auto& ip = ir.IntPoint(q); + trans->SetIntPoint(&ip); + + fe.CalcPhysDShape(*trans, dshape); + double rho_val = m_rho.GetValue(elemID, ip); + double ref_weight = trans->Weight() * ip.weight; + + mfem::DenseMatrix J_map(dim, dim), J_inv(dim, dim); + m_fem.mapping->ComputeJacobian(*trans, J_map); + double detJ = J_map.Det(); + mfem::CalcInverse(J_map, J_inv); + + for (int k = 0; k < dof; ++k) { + for (int c = 0; c < dim; ++c) { + double trace_contrib = 0.0; + for (int j = 0; j < dim; ++j) { + trace_contrib += J_inv(j, c) * dshape(k, j); + } + elvec(k + c * dof) += rho_val * fabs(detJ) * trace_contrib * ref_weight; + } + } + } + m_D.AddElementVector(vdofs, elvec); + } + } + + [[nodiscard]] mfem::Vector& GetVec() const { + return m_D; + } + + void Mult(const mfem::Vector &x, mfem::Vector &y) const override { + if (m_is_col) { + y.SetSize(m_D.Size()); + y = 0.0; + y.Add(x(0), m_D); + } else { + y.SetSize(1); + y(0) = m_D * x; + } + } + + +private: + const FEM& m_fem; + const mfem::GridFunction& m_rho; + const bool m_is_col; + + mutable mfem::Vector m_D; +}; + +template +class ResidualOperator : public mfem::Operator { +public: + ResidualOperator( + FEM& fem, + const Args& args, + const EOS_P& eos_enthalpy, + const EOS_P& eos_pressure, + const double alpha + ) : + Operator(fem.block_offsets.Last()), + m_fem(fem), + m_args(args), + m_eos_enthalpy(eos_enthalpy), + m_eos_pressure(eos_pressure), + m_alpha(std::make_unique(alpha)), + m_fluid_nlf(m_fem.H1_fes.get()), + m_reference_stiffness(m_fem.Vec_H1_fes.get()) + { + m_fluid_nlf.AddDomainIntegrator(new FluidIntegrator(m_eos_enthalpy, m_fem.mapping.get())); + + m_reference_stiffness.AddDomainIntegrator(new mfem::VectorMassIntegrator(*m_alpha)); + m_reference_stiffness.AddDomainIntegrator(new mfem::VectorDiffusionIntegrator()); + m_reference_stiffness.Assemble(); + m_reference_stiffness.Finalize(); + std::println("ResidualOperator initialized with XAD-enabled FluidIntegrator and reference stiffness."); + }; + + + void Mult(const mfem::Vector &u, mfem::Vector &r) const override { + mfem::GridFunction rho, d; + rho.MakeRef(m_fem.H1_fes.get(), u.GetData() + m_fem.block_offsets[0]); + d.MakeRef(m_fem.Vec_H1_fes.get(), u.GetData() + m_fem.block_offsets[1]); + double lambda = u(m_fem.block_offsets[2]); + + m_fem.mapping->SetDisplacement(d); + + m_fem.com = get_com(m_fem, rho); + m_fem.Q = compute_quadrupole_moment_tensor(m_fem, rho, m_fem.com); + + mfem::GridFunction r_rho, r_d; + r_rho.MakeRef(m_fem.H1_fes.get(), r.GetData() + m_fem.block_offsets[0]); + r_d.MakeRef(m_fem.Vec_H1_fes.get(), r.GetData() + m_fem.block_offsets[1]); + + double &r_lambda = r(m_fem.block_offsets[2]); + + m_fluid_nlf.Mult(rho, r_rho); + + auto phi = get_potential(m_fem, m_args, rho); + mfem::GridFunctionCoefficient phi_c(phi.get()); + MappedScalarCoefficient mapped_phi_c(*m_fem.mapping, phi_c); + mfem::LinearForm phi_lf(m_fem.H1_fes.get()); + phi_lf.AddDomainIntegrator(new mfem::DomainLFIntegrator(mapped_phi_c)); + phi_lf.Assemble(); + r_rho += phi_lf; + + mfem::ConstantCoefficient lambda_c(lambda); + MappedScalarCoefficient mapped_lambda_c(*m_fem.mapping, lambda_c); + mfem::LinearForm mass_grad_lf(m_fem.H1_fes.get()); + mass_grad_lf.AddDomainIntegrator(new mfem::DomainLFIntegrator(mapped_lambda_c)); + mass_grad_lf.Assemble(); + + // ReSharper disable once CppDFAUnusedValue + r_rho += mass_grad_lf; + + m_reference_stiffness.Mult(d, r_d); + + PressureBoundaryForce pbf( + m_fem.H1_fes->GetMesh()->Dimension(), + m_fem, + rho, + m_eos_pressure, + m_args.c + ); + mfem::LinearForm pbf_lf(m_fem.Vec_H1_fes.get()); + pbf_lf.AddBoundaryIntegrator(new mfem::VectorBoundaryLFIntegrator(pbf)); + pbf_lf.Assemble(); + + r_d -= pbf_lf; + + for (int i = 0; i < m_fem.ess_tdof_x.Size(); ++i) { + r_d(m_fem.ess_tdof_x[i]) = 0.0; + } + + double current_mass = domain_integrate_grid_function(m_fem, rho); + r_lambda = current_mass - m_args.mass; + } + + [[nodiscard]] Operator& GetGradient(const mfem::Vector &u) const override { + mfem::GridFunction rho, d; + rho.MakeRef(m_fem.H1_fes.get(), u.GetData() + m_fem.block_offsets[0]); + d.MakeRef(m_fem.Vec_H1_fes.get(), u.GetData() + m_fem.block_offsets[1]); + + m_fem.mapping->SetDisplacement(d); + + m_approx_jacobian = std::make_unique(m_fem.block_offsets); + + const mfem::GridFunction grad(m_fem.Vec_H1_fes.get(), u.GetData() + m_fem.block_offsets[1]); + // m_fem.mesh->SetNodes(grad); + + m_approx_jacobian->SetBlock(0, 0, &m_fluid_nlf.GetGradient(rho)); + m_approx_jacobian->SetBlock(1, 1, &m_reference_stiffness); + + B_vec.SetSize(m_fem.H1_fes->GetTrueVSize()); + B_vec = 0.0; + + mfem::ConstantCoefficient one(1.0); + MappedScalarCoefficient mapped_b_vec(*m_fem.mapping, one); + mfem::LinearForm b_lf(m_fem.H1_fes.get()); + b_lf.AddDomainIntegrator(new mfem::DomainLFIntegrator(mapped_b_vec)); + b_lf.Assemble(); + B_vec += b_lf; + + m_B_vec_op_col.SetVector(B_vec, true); + m_B_vec_op_row.SetVector(B_vec, false); + + m_approx_jacobian->SetBlock(0, 2, &m_B_vec_op_col); + m_approx_jacobian->SetBlock(2, 0, &m_B_vec_op_row); + + m_C = std::make_unique>(m_fem, rho, m_eos_pressure); + m_D = std::make_unique>(m_fem, rho, false); + + m_approx_jacobian->SetBlock(1, 0, m_C.get()); + m_approx_jacobian->SetBlock(2, 1, m_D.get()); + + return *m_approx_jacobian; + + } +private: + FEM& m_fem; + const Args& m_args; + const EOS_P& m_eos_enthalpy; + const EOS_P& m_eos_pressure; + const std::unique_ptr m_alpha; + + mutable mfem::NonlinearForm m_fluid_nlf; + mutable mfem::BilinearForm m_reference_stiffness; + mutable std::unique_ptr m_approx_jacobian = nullptr; + mutable mfem::Vector B_vec; + mutable VectorOperator m_B_vec_op_col; + mutable VectorOperator m_B_vec_op_row; + mutable std::unique_ptr> m_C; + mutable std::unique_ptr> m_D; +}; +//endregion + +//region Utility Functions +FEM setup_fem(const std::string& filename, const bool verbose) { + FEM fem; + fem.mesh = std::make_unique(filename, 0, 0); + fem.mesh->EnsureNodes(); + int dim = fem.mesh->Dimension(); + fem.H1_fec = std::make_unique(2, dim); + fem.H1_fes = std::make_unique(fem.mesh.get(), fem.H1_fec.get()); + fem.Vec_H1_fes = std::make_unique(fem.mesh.get(), fem.H1_fec.get(), dim, mfem::Ordering::byNODES); + + auto [r_star_ref, r_inf_ref] = DiscoverBounds(fem.mesh.get(), 3) + .or_else([](const BoundsError& err)->std::expected { + throw std::runtime_error("Unable to determine vacuum domain reference boundary..."); + }).value(); + + fem.mapping = std::make_unique(r_star_ref, r_inf_ref); + + fem.block_offsets.SetSize(4); + fem.block_offsets[0] = 0; + fem.block_offsets[1] = fem.H1_fes->GetTrueVSize(); + fem.block_offsets[2] = fem.block_offsets[1] + fem.Vec_H1_fes->GetTrueVSize(); + fem.block_offsets[3] = fem.block_offsets[2] + 1; + + // Setup Domain Markers to constrain integration + fem.star_marker = std::make_unique>(fem.mesh->attributes.Max()); + fem.star_marker->operator=(0); + fem.star_marker.operator->()[0] = 1; // core + fem.star_marker.operator->()[1] = 1; // envelope + + fem.vacuum_marker = std::make_unique>(fem.mesh->attributes.Max()); + fem.vacuum_marker->operator=(0); + fem.vacuum_marker.operator->()[2] = 1; // vacuum + + fem.surface_marker = std::make_unique>(fem.mesh->bdr_attributes.Max()); + fem.surface_marker->operator=(0); + fem.surface_marker.operator->()[0] = 1; // surface + + fem.vacuum_marker = std::make_unique>(fem.mesh->bdr_attributes.Max()); + fem.vacuum_marker->operator=(0); + fem.vacuum_marker.operator->()[1] = 1; // Infinity + + // Initial COM should be 0, initial mass distribution should be uniform + fem.com.SetSize(dim); fem.com = 0.0; + fem.Q.SetSize(dim, dim); fem.Q = 0.0; + + fem.reference_x = std::make_unique(fem.Vec_H1_fes.get()); + fem.mesh->GetNodes(*fem.reference_x); + + fem.ess_tdof_x.SetSize(0); // No essential boundary conditions for the displacement, the null space here is handled with a weak penalty term + return fem; +} + +void view_mesh(const std::string& host, int port, const mfem::Mesh& mesh, const mfem::GridFunction& gf, const std::string& title) { + mfem::socketstream sol_sock(host.c_str(), port); + if (!sol_sock.is_open()) return; + sol_sock << "solution\n" << mesh << gf; + sol_sock << "window_title '" << title << "'\n" << std::flush; +} + +double domain_integrate_grid_function(const FEM& fem, const mfem::GridFunction& gf, Domains domain) { + mfem::LinearForm lf(fem.H1_fes.get()); + mfem::GridFunctionCoefficient gf_c(&gf); + double integral; + mfem::Array elem_markers; + populate_element_mask(fem, domain, elem_markers); + + if (fem.has_mapping()) { + MappedScalarCoefficient mapped_gf_c(*fem.mapping, gf_c); + lf.AddDomainIntegrator(new mfem::DomainLFIntegrator(mapped_gf_c), elem_markers); + lf.Assemble(); + integral = lf.Sum(); + } else { + lf.AddDomainIntegrator(new mfem::DomainLFIntegrator(gf_c), elem_markers); + lf.Assemble(); + integral = lf.Sum(); + } + return integral; +} + +mfem::Vector get_com(const FEM& fem, const mfem::GridFunction &rho) { + const int dim = fem.mesh->Dimension(); + mfem::Vector com(dim); + com = 0.0; + double total_mass = 0.0; + + for (int i = 0; i < fem.H1_fes->GetNE(); ++i) { + mfem::ElementTransformation *trans = fem.H1_fes->GetElementTransformation(i); + const mfem::IntegrationRule &ir = mfem::IntRules.Get(trans->GetGeometryType(), fem.H1_fes->GetOrder(0) + trans->OrderW()); + + for (int j = 0; j < ir.GetNPoints(); ++j) { + const mfem::IntegrationPoint &ip = ir.IntPoint(j); + trans->SetIntPoint(&ip); + + double weight = trans->Weight() * ip.weight; + if (fem.has_mapping()) { + weight *= fem.mapping->ComputeDetJ(*trans, ip); + } + double rho_val = rho.GetValue(i, ip); + + mfem::Vector phys_point(dim); + if (fem.has_mapping()) { + fem.mapping->GetPhysicalPoint(*trans, ip, phys_point); + } else { + trans->Transform(ip, phys_point); + } + + const double mass_term = rho_val * weight; + total_mass += mass_term; + + for (int d = 0; d < dim; ++d) { + com(d) += phys_point(d) * mass_term; + } + } + } + com /= total_mass; + return com; +} + +void get_physical_coordinates(const mfem::GridFunction& reference_pos, const mfem::GridFunction& displacement, mfem::GridFunction& physical_pos) { + add(reference_pos, displacement, physical_pos); +} + +void populate_element_mask(const FEM &fem, Domains domain, mfem::Array &mask) { + mask.SetSize(fem.mesh->attributes.Max()); + mask = 0; + + if ((domain & Domains::CORE) == Domains::CORE) { + mask[0] = 1; + } + + if ((domain & Domains::ENVELOPE) == Domains::ENVELOPE) { + mask[1] = 1; + } + + if ((domain & Domains::VACUUM) == Domains::VACUUM) { + mask[2] = 1; + } +} + +std::expected DiscoverBounds(const mfem::Mesh *mesh, const int vacuum_attr) { + double min_r = std::numeric_limits::max(); + double max_r = -std::numeric_limits::max(); + + bool found_vacuum = false; + + for (int i = 0; i < mesh->GetNE(); ++i) { + if (mesh->GetAttribute(i) == vacuum_attr) { + found_vacuum = true; + mfem::Array vertices; + mesh->GetElementVertices(i, vertices); + for (const int v: vertices) { + const double* coords = mesh->GetVertex(v); + double r = std::sqrt(coords[0] * coords[0] + coords[1] * coords[1] + coords[2] * coords[2]); + min_r = std::min(min_r, r); + max_r = std::max(max_r, r); + } + } + } + if (found_vacuum) { + return Bounds(min_r, max_r); + } + return std::unexpected(BoundsError::CANNOT_FIND_VACUUM); +} + +void conserve_mass(const FEM& fem, mfem::GridFunction& rho, const double target_mass) { + if (const double current_mass = domain_integrate_grid_function(fem, rho); current_mass > 1e-15) rho *= (target_mass / current_mass); +} +//endregion + +//region Physics Functions +double centrifugal_potential(const mfem::Vector& phys_x, double omega) { + const double s2 = std::pow(phys_x(0), 2) + std::pow(phys_x(1), 2); + return -0.5 * s2 * std::pow(omega, 2); +} + +double get_moment_of_inertia(const FEM& fem, const mfem::GridFunction& rho) { + auto s2_func = [](const mfem::Vector& x) { + return std::pow(x(0), 2) + std::pow(x(1), 2); + }; + + std::unique_ptr s2_coeff; + if (fem.has_mapping()) { + s2_coeff = std::make_unique(*fem.mapping, s2_func); + } else { + s2_coeff = std::make_unique(s2_func); + } + + mfem::GridFunctionCoefficient rho_coeff(&rho); + mfem::ProductCoefficient I_integrand(rho_coeff, *s2_coeff); + + mfem::LinearForm I_lf(fem.H1_fes.get()); + + double I = 0.0; + if (fem.has_mapping()) { + MappedScalarCoefficient mapped_integrand(*fem.mapping, I_integrand); + I_lf.AddDomainIntegrator(new mfem::DomainLFIntegrator(mapped_integrand)); + I_lf.Assemble(); + I = I_lf.Sum(); + } else { + I_lf.AddDomainIntegrator(new mfem::DomainLFIntegrator(I_integrand)); + I_lf.Assemble(); + I = I_lf.Sum(); + } + + return I; +} +//endregion + +//region Analytic External Potential for an oblate spheroid (doi.10.16.j.pss.2018.01.005) + +inline double sq(const double x) { return x * x; } + +double eccentricity(const double a, const double c) { + return std::sqrt(1.0 - sq(c)/sq(a)); +} + +double equatorial_radius(const mfem::Vector& x) { + return std::sqrt(sq(x(0)) + sq(x(1))); +} + +double beta_ext(const mfem::Vector& x, const double a, const double c) { + const double e = eccentricity(a, c); + const double z = x(2); + const double r = equatorial_radius(x); + const double ae_sq = sq(a*e); + + if (std::abs(r) < 1e-12) { + return std::atan2(a*e, std::abs(x(2))); + } + + const double fraction = (sq(z) + ae_sq) / sq(r); + const double first_term = -fraction; + + const double radical_one = sq(1.0 + fraction); + const double radical_two = 4.0 * ae_sq / sq(r); + const double second_term = std::sqrt(std::max(0.0, radical_one - radical_two)); + + const double cos_term = first_term + second_term; + return 0.5 * std::acos(std::clamp(cos_term, -1.0, 1.0)); +} + +double oblate_spheroid_surface_potential(const mfem::Vector& x, const double a, const double c, const double total_mass) { + const double beta = beta_ext(x, a, c); + const double r = equatorial_radius(x); + const double e = eccentricity(a, c); + const double z = x(2); + + const double G_M = G * total_mass; + const double e3a3 = std::pow(e, 3) * std::pow(a, 3); + + const double t1 = (3.0 * G_M * beta) / (2.0 * e * a); + + const double t2aa = (3.0 * G_M * sq(r)) / (2.0 * e3a3); + const double t2ab = beta - std::sin(beta) * std::cos(beta); + + const double t2ba = (3.0 * G_M * sq(z)) / e3a3; + const double t2bb = std::tan(beta) - beta; + + return -t1 + 0.5 * (t2aa * t2ab + t2ba * t2bb); +} +//endregion End Analytic Potential + +//region potentials +std::unique_ptr grav_potential(const FEM& fem, const Args& args, const mfem::GridFunction& rho) { + auto phi = std::make_unique(fem.H1_fes.get()); + + mfem::Array ess_bdr(fem.mesh->bdr_attributes.Max()); + ess_bdr = 0; + assert(ess_bdr.Size() == 2); + ess_bdr[1] = 1; // The only boundary condition is that the potential goes to 0 at infinity + + // mfem::ConstantCoefficient zero(0.0); + // phi->ProjectBdrCoefficient(zero, ess_bdr); // Set the potential to 0 at infinity; + + double total_mass = domain_integrate_grid_function(fem, rho, STELLAR); + auto bdr_func = [&fem, total_mass](const mfem::Vector& x) { + double r = x.Norml2(); + double theta = std::atan2(x(1), x(0)); + double phi = std::acos(x(2) / r); + double val = l2_multipole_potential(fem, total_mass, x); + // std::println("{},{},{},{},{}", r, theta, phi, val, -G*total_mass / r); + return l2_multipole_potential(fem, total_mass, x); + }; + + std::unique_ptr phi_bdr_coeff; + if (fem.has_mapping()) { + phi_bdr_coeff = std::make_unique(*fem.mapping, bdr_func); + } else { + phi_bdr_coeff = std::make_unique(bdr_func); + } + phi->ProjectBdrCoefficient(*phi_bdr_coeff, ess_bdr); + + + mfem::Array ess_tdof_list; + fem.H1_fes->GetEssentialTrueDofs(ess_bdr, ess_tdof_list); + + mfem::Array stellar_marker; + populate_element_mask(fem, STELLAR, stellar_marker); + + auto laplacian = std::make_unique(fem.H1_fes.get()); + + mfem::ConstantCoefficient one_coeff(1.0); + std::unique_ptr mapped_diff_coeff; + if (fem.has_mapping()) { + mapped_diff_coeff = std::make_unique(*fem.mapping, one_coeff, fem.mesh->Dimension()); + laplacian->AddDomainIntegrator(new mfem::DiffusionIntegrator(*mapped_diff_coeff)); // The laplacian is global + } else { + laplacian->AddDomainIntegrator(new mfem::DiffusionIntegrator()); + } + + laplacian->Assemble(); + laplacian->Finalize(); + + mfem::LinearForm b(fem.H1_fes.get()); + mfem::ConstantCoefficient four_pi_G(-4.0 * M_PI * G); + mfem::GridFunctionCoefficient rho_coeff(&rho); + mfem::ProductCoefficient rhs_coeff(rho_coeff, four_pi_G); + + auto mapped_rhs = std::make_unique(*fem.mapping, rhs_coeff, MappedScalarCoefficient::EVAL_POINTS::REFERENCE); + if (fem.has_mapping()) { + b.AddDomainIntegrator(new mfem::DomainLFIntegrator(*mapped_rhs), stellar_marker); // The mass contribution is local to the stellar domain + } else { + b.AddDomainIntegrator(new mfem::DomainLFIntegrator(rhs_coeff), stellar_marker); + } + + b.Assemble(); + + mfem::OperatorPtr A; + mfem::Vector B, X; + laplacian->FormLinearSystem(ess_tdof_list, *phi, b, A, X, B); + + mfem::GSSmoother prec; + mfem::CGSolver cg; + cg.SetPreconditioner(prec); + cg.SetOperator(*A); + cg.SetRelTol(args.p.tol); + cg.SetMaxIter(args.p.max_iters); + cg.Mult(B, X); + + laplacian->RecoverFEMSolution(X, b, *phi); + return phi; +} + +// std::unique_ptr grav_potential(const FEM& fem, const Args &args, const mfem::GridFunction& rho, std::optional oblate) { +// auto phi = std::make_unique(fem.H1_fes.get()); +// +// mfem::Array ess_bdr(fem.mesh->bdr_attributes.Max()); +// ess_bdr = 1; +// +// mfem::GridFunctionCoefficient rho_coeff(&rho); +// double total_mass = domain_integrate_grid_function(fem, rho); +// +// auto grav_potential = [&fem, &total_mass, &oblate](const mfem::Vector& x) { +// if (!oblate.has_value() or oblate->use == false) { +// return l2_multipole_potential(fem, total_mass, x); +// } +// return oblate_spheroid_surface_potential(x, oblate->a, oblate->c, total_mass); +// }; +// +// std::unique_ptr phi_bdr_coeff; +// if (fem.has_mapping()) { +// phi_bdr_coeff = std::make_unique(*fem.mapping, grav_potential); +// } else { +// phi_bdr_coeff = std::make_unique(grav_potential); +// } +// phi->ProjectBdrCoefficient(*phi_bdr_coeff, ess_bdr); +// +// mfem::Array ess_tdof_list; +// fem.H1_fes->GetEssentialTrueDofs(ess_bdr, ess_tdof_list); +// +// auto laplacian = std::make_unique(fem.H1_fes.get()); +// // ReSharper disable once CppTooWideScope +// std::unique_ptr one_coeff; +// // ReSharper disable once CppTooWideScope +// std::unique_ptr mapped_diff_coeff; +// if (fem.has_mapping()) { +// one_coeff = std::make_unique(1.0); +// mapped_diff_coeff = std::make_unique(*fem.mapping, *one_coeff, fem.mesh->Dimension()); +// laplacian->AddDomainIntegrator(new mfem::DiffusionIntegrator(*mapped_diff_coeff)); +// } else { +// laplacian->AddDomainIntegrator(new mfem::DiffusionIntegrator()); +// } +// laplacian->Assemble(); +// laplacian->Finalize(); +// +// mfem::ConstantCoefficient four_pi_G(-4.0 * M_PI * G); +// mfem::ProductCoefficient rhs_coeff(rho_coeff, four_pi_G); +// mfem::LinearForm b(fem.H1_fes.get()); +// +// // ReSharper disable once CppTooWideScope +// std::unique_ptr mapped_rhs; +// if (fem.has_mapping()) { +// mapped_rhs = std::make_unique(*fem.mapping, rhs_coeff, MappedScalarCoefficient::EVAL_POINTS::REFERENCE); +// b.AddDomainIntegrator(new mfem::DomainLFIntegrator(*mapped_rhs)); +// } else { +// b.AddDomainIntegrator(new mfem::DomainLFIntegrator(rhs_coeff)); +// } +// b.Assemble(); +// +// mfem::OperatorPtr A; +// mfem::Vector B, X; +// laplacian->FormLinearSystem(ess_tdof_list, *phi, b, A, X, B); +// +// mfem::GSSmoother prec; +// mfem::CGSolver cg; +// cg.SetPreconditioner(prec); +// cg.SetOperator(*A); +// cg.SetRelTol(args.p.tol); +// cg.SetMaxIter(args.p.max_iters); +// cg.SetPrintLevel(0); +// +// cg.Mult(B, X); +// laplacian->RecoverFEMSolution(X, b, *phi); +// return phi; +// } + +std::unique_ptr get_potential(const FEM &fem, const Args &args, const mfem::GridFunction &rho) { + auto phi = grav_potential(fem, args, rho); + + if (args.r.enabled) { + auto rot = [&args](const mfem::Vector& x) { + return centrifugal_potential(x, args.r.omega); + }; + + std::unique_ptr centrifugal_coeff; + if (fem.has_mapping()) { + centrifugal_coeff = std::make_unique(*fem.mapping, rot); + } else { + centrifugal_coeff = std::make_unique(rot); + } + mfem::GridFunction centrifugal_gf(fem.H1_fes.get()); + centrifugal_gf.ProjectCoefficient(*centrifugal_coeff); + (*phi) += centrifugal_gf; + } + + return phi; +} + +mfem::DenseMatrix compute_quadrupole_moment_tensor(const FEM& fem, const mfem::GridFunction& rho, const mfem::Vector& com) { + const int dim = fem.mesh->Dimension(); + mfem::DenseMatrix Q(dim, dim); + Q = 0.0; + + for (int i = 0; i < fem.H1_fes->GetNE(); ++i) { + mfem::ElementTransformation *trans = fem.mesh->GetElementTransformation(i); + const mfem::IntegrationRule &ir = mfem::IntRules.Get(trans->GetGeometryType(), 2 * fem.H1_fes->GetOrder(0) + trans->OrderW()); + for (int j = 0; j < ir.GetNPoints(); ++j) { + const mfem::IntegrationPoint &ip = ir.IntPoint(j); + trans->SetIntPoint(&ip); + + double weight = trans->Weight() * ip.weight; + + if (fem.has_mapping()) { + weight *= fem.mapping->ComputeDetJ(*trans, ip); + } + + const double rho_val = rho.GetValue(i, ip); + + mfem::Vector phys_point(dim); + if (fem.has_mapping()) { + fem.mapping->GetPhysicalPoint(*trans, ip, phys_point); + } else { + trans->Transform(ip, phys_point); + } + + mfem::Vector x_prime(dim); + double r_sq = 0.0; + + for (int d = 0; d < dim; ++d) { + x_prime(d) = phys_point(d) - com(d); + r_sq += x_prime(d) * x_prime(d); + } + + for (int m = 0; m < dim; ++m) { + for (int n = 0; n < dim; ++n) { + const double delta = (m == n) ? 1.0 : 0.0; + const double contrib = 3.0 * x_prime(m) * x_prime(n) - delta * r_sq; + Q(m, n) += rho_val * contrib * weight; + } + } + } + } + return Q; +} + +double l2_multipole_potential(const FEM &fem, const double total_mass, const mfem::Vector &phys_x) { + const double r = phys_x.Norml2(); + if (r < 1e-12) return 0.0; + + const int dim = fem.mesh->Dimension(); + + mfem::Vector n(phys_x); + n /= r; + + double l2_mult_factor = 0.0; + for (int i = 0; i < dim; ++i) { + for (int j = 0; j < dim; ++j) { + l2_mult_factor += fem.Q(i, j) * n(i) * n(j); + } + } + + const double l2_contrib = (G / (2.0 * std::pow(r, 3))) * l2_mult_factor; + + const double l0_contrib = -G * total_mass / r; + + // l1 contribution is zero for a system centered on its COM + return l0_contrib + l2_contrib; +} +//endregion + +//region Tests +void test_mesh_load(const FEM& fem) { + size_t failed = 0; + if (not fem.okay()) ++failed; + const int dim = fem.mesh->Dimension(); + if (dim != 3) ++failed; + + const int n_scalar = fem.H1_fes->GetTrueVSize(); + const int n_vector = fem.Vec_H1_fes->GetTrueVSize(); + if (n_vector != dim * n_scalar) ++failed; + + if (fem.block_offsets[0] != 0) ++failed; + if (fem.block_offsets[1] != n_scalar) ++failed; + if (fem.block_offsets[2] != n_scalar + n_vector) ++failed; + if (fem.block_offsets[3] != n_scalar + n_vector + 1) ++failed; + + constexpr size_t num_tests = 6; + auto result_type = TEST_RESULT_TYPE::FAILURE; + if (failed == 0) { + result_type = TEST_RESULT_TYPE::SUCCESS; + } else if (failed < num_tests) { + result_type = TEST_RESULT_TYPE::PARTIAL; + } + std::string test_msg = fmt_test_msg("Mesh Load Test", result_type, failed, num_tests); + std::println("{}", test_msg); + + assert(dim == 3); + assert(n_vector == (n_scalar * dim)); + assert (fem.block_offsets[0] == 0); + assert (fem.block_offsets[1] == n_scalar); + assert (fem.block_offsets[2] == n_scalar + n_vector); + assert (fem.block_offsets[3] == n_scalar + n_vector + 1); +} + +void test_ref_coord_storage(const FEM& fem) { + size_t failed = 0; + if (not fem.mapping->IsIdentity()) ++failed; + + const size_t num_elemIDs = std::min(30, fem.mesh->GetNE()); + for (int elemID = 0; elemID < num_elemIDs; ++elemID) { + auto* trans = fem.mesh->GetElementTransformation(elemID); + const auto& ir = mfem::IntRules.Get(trans->GetGeometryType(), 2); + const auto& ip = ir.IntPoint(0); + trans->SetIntPoint(&ip); + + mfem::Vector x_ref, x_phys; + trans->Transform(ip, x_ref); + fem.mapping->GetPhysicalPoint(*trans, ip, x_phys); + x_ref -= x_phys; + + if (not (x_ref.Norml2() < 1e-12)) ++failed; + } + + const size_t num_tests = num_elemIDs + 1; + auto result_type = TEST_RESULT_TYPE::FAILURE; + if (failed == 0) { + result_type = TEST_RESULT_TYPE::SUCCESS; + } else if (failed < num_tests) { + result_type = TEST_RESULT_TYPE::PARTIAL; + } + + std::string test_msg = fmt_test_msg("Mesh Ref Coord", result_type, failed, num_tests); + std::println("{}", test_msg); +} + +void test_reference_volume_integral(const FEM& fem) { + size_t failed = 0; + + mfem::GridFunction ones(fem.H1_fes.get()); + ones = 1.0; + + double vol = domain_integrate_grid_function(fem, ones, STELLAR); + double expected = (4.0/3.0) * M_PI * std::pow(RADIUS, 3.0); + double rel_err = std::abs(vol - expected) / expected; + + if (rel_err > 1e-2) ++failed; + + constexpr size_t num_tests = 1; + auto result_type = TEST_RESULT_TYPE::FAILURE; + if (failed == 0) { + result_type = TEST_RESULT_TYPE::SUCCESS; + } + + std::println("{}", fmt_test_msg("Reference Volume Integral", result_type, failed, num_tests)); + + if (result_type == TEST_RESULT_TYPE::FAILURE) { + std::println("\tFAILURE: Volume: {}, Expected: {}, Error (rel): {}", vol, expected, rel_err); + } +} + +void test_spherically_symmetric_com(const FEM& fem) { + mfem::GridFunction rho(fem.H1_fes.get()); + rho = 1.0; + + mfem::Vector com = get_com(fem, rho); + size_t failed = 0; + + const size_t dim = fem.mesh->Dimension(); + const size_t num_tests = dim; + + for (int dimID = 0; dimID < num_tests; ++dimID) { + if (std::abs(com(dimID)) > 1e-12) ++failed; + } + + auto result_type = TEST_RESULT_TYPE::FAILURE; + if (failed == 0) { + result_type = TEST_RESULT_TYPE::SUCCESS; + } else if (failed < num_tests) { + result_type = TEST_RESULT_TYPE::PARTIAL; + } + + std::println("{}", fmt_test_msg("Uniform COM", result_type, failed, num_tests)); + + if (result_type == TEST_RESULT_TYPE::FAILURE) { + std::println("\t COM=<{:+0.3E}, {:+0.3E}, {:+0.3E}>", com(0), com(1), com(2)); + } +} + +void test_com_variance_to_displacement(const FEM& fem) { + size_t failed = 0; + mfem::GridFunction linear_displacement(fem.Vec_H1_fes.get()); + linear_displacement = 10.0; // This will linearly displace the domain by 10 unit along all axes + + fem.mapping->SetDisplacement(linear_displacement); + + mfem::GridFunction rho(fem.H1_fes.get()); + rho = 1.0; + + mfem::Vector mapped_com = get_com(fem, rho); + + const size_t dim = fem.mesh->Dimension(); + const size_t num_tests = dim; + for (int dimID = 0; dimID < num_tests; ++dimID) { + if (10 - std::abs(mapped_com(dimID)) > 1e-12) ++failed; + } + + auto result_type = TEST_RESULT_TYPE::FAILURE; + if (failed == 0) { + result_type = TEST_RESULT_TYPE::SUCCESS; + } else if (failed < num_tests) { + result_type = TEST_RESULT_TYPE::PARTIAL; + } + + std::println("{}", fmt_test_msg("COM variance to displacement", result_type, failed, num_tests)); + + if (result_type == TEST_RESULT_TYPE::FAILURE) { + std::println("\tFAILURE COM=<{:+0.2E}, {:+0.2E}, {:+0.2E}>", mapped_com(0), mapped_com(1), mapped_com(2)); + } + + fem.mapping->ResetDisplacement(); +} + +void test_volume_invariance_to_displacement(const FEM& fem) { + size_t failed = 0; + mfem::GridFunction linear_displacement(fem.Vec_H1_fes.get()); + linear_displacement = 10.0; // This will linearly displace the domain by 10 unit along all axes + + fem.mapping->SetDisplacement(linear_displacement); + + mfem::GridFunction ones(fem.H1_fes.get()); + ones = 1.0; + double mapped_vol = domain_integrate_grid_function(fem, ones, STELLAR); + double expected = (4.0/3.0) * M_PI * std::pow(RADIUS, 3.0); + double rel_err = std::abs(mapped_vol - expected) / expected; + + if (rel_err > 1e-2) ++failed; + constexpr size_t num_tests = 1; + auto result_type = TEST_RESULT_TYPE::FAILURE; + if (failed == 0) { + result_type = TEST_RESULT_TYPE::SUCCESS; + } + + std::println("{}", fmt_test_msg("Invariance of volume against translation", result_type, failed, num_tests)); + + if (result_type == TEST_RESULT_TYPE::FAILURE) { + std::println("\tFAILURE: Volume: {}, Expected: {}", mapped_vol, expected); + } + fem.mapping->ResetDisplacement(); +} + +void test_volume_ellipsoid_deformation(const FEM& fem) { + size_t failed = 0; + size_t num_tests = 0; + + constexpr double a = 2.0; // x-axis + constexpr double b = 0.5; // y-axis + constexpr double c = 1.5; // z-axis + constexpr double expected_vol = (4.0 / 3.0) * M_PI * a * b * c; + + mfem::GridFunction ellipsoid_displacement(fem.Vec_H1_fes.get()); + { + const int dim = fem.mesh->Dimension(); + mfem::VectorFunctionCoefficient disp_coeff(dim, [&](const mfem::Vector& x, mfem::Vector& d) { + d.SetSize(x.Size()); + d(0) = (a - 1.0) * x(0); + d(1) = (b - 1.0) * x(1); + d(2) = (c - 1.0) * x(2); + }); + ellipsoid_displacement.ProjectCoefficient(disp_coeff); + } + fem.mapping->SetDisplacement(ellipsoid_displacement); + + { + ++num_tests; + mfem::GridFunction ones(fem.H1_fes.get()); + ones = 1.0; + const double mapped_vol = domain_integrate_grid_function(fem, ones, STELLAR); + const double rel_err = std::abs(mapped_vol - expected_vol) / expected_vol; + if (rel_err > 1e-3) { + ++failed; + std::println("\tFAILURE Test 1: Mapped volume = {:.6f}, expected = {:.6f}, rel_err = {:.2e}", + mapped_vol, expected_vol, rel_err); + } + } + + { + ++num_tests; + const double expected_x2_integral = std::pow(a, 3) * b * c * (4.0 * M_PI / 15.0); + + mfem::GridFunction x_ref_sq(fem.H1_fes.get()); + mfem::FunctionCoefficient x_sq_coeff([](const mfem::Vector& x) { + return x(0) * x(0); + }); + x_ref_sq.ProjectCoefficient(x_sq_coeff); + + mfem::GridFunction x_phys_sq(fem.H1_fes.get()); + PhysicalPositionFunctionCoefficient x_phys_sq_coeff(*fem.mapping, + [](const mfem::Vector& x_phys) { + return x_phys(0) * x_phys(0); + } + ); + x_phys_sq.ProjectCoefficient(x_phys_sq_coeff); + + const double mapped_x2_integral = domain_integrate_grid_function(fem, x_phys_sq, STELLAR); + const double rel_err = std::abs(mapped_x2_integral - expected_x2_integral) / expected_x2_integral; + if (rel_err > 1e-3) { + ++failed; + std::println("\tFAILURE Test 2: integral x_phys^2 = {:.6f}, expected = {:.6f}, rel_err = {:.2e}", + mapped_x2_integral, expected_x2_integral, rel_err); + } + } + + { + ++num_tests; + constexpr double expected_detJ = a * b * c; + double max_detJ_err = 0.0; + for (int e = 0; e < std::min(5, fem.mesh->GetNE()); ++e) { + auto* trans = fem.mesh->GetElementTransformation(e); + const auto& ir = mfem::IntRules.Get(trans->GetGeometryType(), 2); + for (int q = 0; q < ir.GetNPoints(); ++q) { + const auto& ip = ir.IntPoint(q); + trans->SetIntPoint(&ip); + const double detJ = fem.mapping->ComputeDetJ(*trans, ip); + max_detJ_err = std::max(max_detJ_err, std::abs(detJ - expected_detJ)); + } + } + if (max_detJ_err > 1e-10) { + ++failed; + std::println("\tFAILURE Test 3: max pointwise |det(J) - a*b*c| = {:.2e}", max_detJ_err); + } + } + + auto result_type = TEST_RESULT_TYPE::FAILURE; + if (failed == 0) { + result_type = TEST_RESULT_TYPE::SUCCESS; + } else if (failed < num_tests) { + result_type = TEST_RESULT_TYPE::PARTIAL; + } + + std::println("{}", fmt_test_msg("Volume under ellipsoidal deformation", result_type, failed, num_tests)); + + fem.mapping->ResetDisplacement(); +} + +void test_uniform_potential(FEM& fem, const Args& args) { + const double analytic_vol = (4.0/3.0) * M_PI * std::pow(RADIUS, 3); + const double rho0 = MASS / analytic_vol; + + mfem::GridFunction rho_uniform(fem.H1_fes.get()); + rho_uniform = rho0; + + fem.com = get_com(fem, rho_uniform); + fem.Q = compute_quadrupole_moment_tensor(fem, rho_uniform, fem.com); + + const auto phi = grav_potential(fem, args, rho_uniform); + + double max_abs_err = 0.0; + double max_rel_err = 0.0; + constexpr double tol = 1e-3; + + size_t failed = 0; + size_t num_tests = 0; + const size_t num_elemIDs = std::min(30, fem.mesh->GetNE()); + for (int elemID = 0; elemID < num_elemIDs; ++elemID) { + num_tests++; + auto* trans = fem.mesh->GetElementTransformation(elemID); + const auto& ir = mfem::IntRules.Get(trans->GetGeometryType(), 2); + const auto& ip = ir.IntPoint(0); + trans->SetIntPoint(&ip); + + mfem::Vector x_ref, x_phys; + trans->Transform(ip, x_ref); + fem.mapping->GetPhysicalPoint(*trans, ip, x_phys); + + const double r = x_phys.Norml2(); + + if (r < 1e-9) continue; + + const double phi_analytic = (-G * MASS / (2.0 * std::pow(RADIUS, 3.0))) * (3.0*RADIUS*RADIUS - r*r); + + const double phi_fem = phi->GetValue(elemID, ip); + + const double abs_err = std::abs(phi_fem - phi_analytic); + const double rel_err = abs_err / std::abs(phi_analytic); + + max_abs_err = std::max(max_abs_err, abs_err); + max_rel_err = std::max(max_rel_err, rel_err); + + if (rel_err > tol) ++failed; + } + + auto result_type = TEST_RESULT_TYPE::FAILURE; + if (failed == 0) { + result_type = TEST_RESULT_TYPE::SUCCESS; + } else if (failed < num_tests) { + result_type = TEST_RESULT_TYPE::PARTIAL; + } + + std::println("{}", fmt_test_msg("Test Uniform Potential", result_type, failed, num_tests)); + + if (result_type == TEST_RESULT_TYPE::FAILURE) { + std::println("\tFAILURE: max abs error: {:+0.2E}, max rel error: {:+0.2E}", max_abs_err, max_rel_err); + } +} + +void test_ellipsoidal_potential(FEM& fem, const Args& args) { + constexpr double a = 1.0 * RADIUS; + constexpr double b = a; // oblate + constexpr double c = 0.99 * RADIUS; + + constexpr double expected_vol = (4.0 / 3.0) * M_PI * a * b * c; + constexpr double rho0 = MASS / expected_vol; + + mfem::GridFunction ellipsoidal_disp(fem.Vec_H1_fes.get()); + mfem::VectorFunctionCoefficient disp_coeff(3, [&](const mfem::Vector& x, mfem::Vector& d) { + d.SetSize(3); + d(0) = (a/RADIUS - 1.0) * x(0); + d(1) = (b/RADIUS - 1.0) * x(1); + d(2) = (c/RADIUS - 1.0) * x(2); + }); + ellipsoidal_disp.ProjectCoefficient(disp_coeff); + fem.mapping->SetDisplacement(ellipsoidal_disp); + + mfem::GridFunction rho(fem.H1_fes.get()); + rho = rho0; + + fem.com = get_com(fem, rho); + fem.Q = compute_quadrupole_moment_tensor(fem, rho, fem.com); + + // OblatePotential oblate{.use=true, .a=a, .c=c,.rho_0=rho0}; + const auto phi = grav_potential(fem, args, rho); + + constexpr double e_sq = 1.0 - (c * c)/(a*a); + const double e = std::sqrt(e_sq); + + const double I_const = (2.0 * std::sqrt(1.0 - e_sq) / e) * std::asin(e); + const double A_R = (std::sqrt(1.0-e_sq) / std::pow(e, 3.0)) * std::asin(e) - (1.0 - e_sq)/e_sq; + const double A_z = (2.0 / e_sq) * (1.0 - (std::sqrt(1.0-e_sq) / e) * std::asin(e)); + + size_t failed = 0; + size_t num_tests = 0; + double max_rel_err = 0.0; + double total_err = 0.0; + const size_t check_count = std::min(50, fem.mesh->GetNE()); + + for (int elemID = 0; elemID < check_count; ++elemID) { + auto* trans = fem.mesh->GetElementTransformation(elemID); + const auto& ip = mfem::IntRules.Get(trans->GetGeometryType(), 2).IntPoint(0); + trans->SetIntPoint(&ip); + + mfem::Vector x_phys; + fem.mapping->GetPhysicalPoint(*trans, ip, x_phys); + + const double R2 = x_phys(0)*x_phys(0) + x_phys(1)*x_phys(1); + const double z2 = x_phys(2)*x_phys(2); + const double phi_analytic = -M_PI * G * rho0 * (a*a*I_const - A_R * R2 - A_z * z2); + + const double phi_fem = phi->GetValue(elemID, ip); + const double rel_err = std::abs(phi_fem - phi_analytic) / std::abs(phi_analytic); + max_rel_err = std::max(max_rel_err, rel_err); + total_err += rel_err; + num_tests++; + if (rel_err > 1e-3) ++failed; + } + + auto result_type = TEST_RESULT_TYPE::FAILURE; + if (failed == 0) { + result_type = TEST_RESULT_TYPE::SUCCESS; + } else if (failed < num_tests) { + result_type = TEST_RESULT_TYPE::PARTIAL; + } + + std::println("{}", fmt_test_msg("Test Ellipsoidal Potential", result_type, failed, num_tests)); + if (result_type == TEST_RESULT_TYPE::FAILURE) { + std::println("\tFAILURE: max rel error: {:+0.2E}, mean rel error: {:+0.2E}", max_rel_err, total_err/static_cast(num_tests)); + } + +} +//endregion + +int main(int argc, char** argv) { + Args args; + + CLI::App app{"Mapped Coordinates XAD-Enabled Non-Linear Solver"}; + app.add_option("-m,--mesh", args.mesh_file)->required()->check(CLI::ExistingFile); + app.add_option("--max-iters", args.max_iters)->default_val(20); + app.add_option("--index", args.index)->default_val(1); + app.add_option("--mass", args.mass)->default_val(MASS); + app.add_option("--surface-pressure-scalar", args.c)->default_val(1e-4); + + args.r.enabled = false; + args.p.tol = 1e-6; + args.p.max_iters = 1000; + + CLI11_PARSE(app, argc, argv); + FEM fem = setup_fem(args.mesh_file, true); + // fem.mesh->UniformRefinement(); + + + test_mesh_load(fem); + test_ref_coord_storage(fem); + test_reference_volume_integral(fem); + test_spherically_symmetric_com(fem); + + test_com_variance_to_displacement(fem); + test_volume_invariance_to_displacement(fem); + test_volume_ellipsoid_deformation(fem); + // test_kelvin_jacobian(fem); + test_uniform_potential(fem, args); + test_ellipsoidal_potential(fem, args); // Note that this test is currently predicated on an analytic solution for the surface boundary potential + + CoupledState state(fem); + + // typedef xad::AReal adouble; + // + // EOS_P eos_enthalpy = [index = args.index](const adouble& rho, const adouble& temp) -> adouble { + // if (rho.value() < 1e-15) return {0.0}; + // return (index + 1.0) * pow(rho, 1.0 / index); + // }; + // + // EOS_P eos_pressure = [index = args.index](const adouble& rho, const adouble& temp) -> adouble { + // if (rho.value() < 1e-15) return {0.0}; + // return pow(rho, 1.0 + (1.0 / index)); + // }; + // + // auto init_rho_func = [&](const mfem::Vector& pt) { + // const double r = pt.Norml2(); + // return (r < RADIUS) ? (1.0 - r/RADIUS) : 0.0; + // }; + // mfem::FunctionCoefficient init_rho_coeff(init_rho_func); + // state.rho.ProjectCoefficient(init_rho_coeff); + // conserve_mass(fem, state.rho, args.mass); + // view_mesh(HOST, PORT, *fem.mesh, state.rho, "Initial Density Solution with XAD AD"); + // view_mesh(HOST, PORT, *fem.mesh, state.d, "Initial Position Solution with XAD AD"); + // + // std::println("Starting Coupled Block Solver with XAD AD..."); + // + // ResidualOperator coupled_operator(fem, args, eos_enthalpy, eos_pressure); + // + // mfem::NewtonSolver newton_solver; + // newton_solver.SetOperator(coupled_operator); + // newton_solver.SetMaxIter(500); + // newton_solver.SetRelTol(1e-6); + // newton_solver.SetPrintLevel(1); + // + // mfem::GMRESSolver jfnk_solver; + // jfnk_solver.SetMaxIter(100); + // jfnk_solver.SetRelTol(1e-4); + // newton_solver.SetSolver(jfnk_solver); + // newton_solver.SetAdaptiveLinRtol(2, 0.5, 0.9, 0.9, 1.1); + // + // mfem::Vector zero_rhs(fem.block_offsets.Last()); + // zero_rhs = 0.0; + // newton_solver.Mult(zero_rhs, *state.U); + // + // mfem::GridFunction rho(fem.H1_fes.get(), state.U->GetBlock(0), 0); + // mfem::GridFunction x(fem.Vec_H1_fes.get(), state.U->GetBlock(1), 0); + // view_mesh(HOST, PORT, *fem.mesh, rho, "Final Density Solution with XAD AD"); + // view_mesh(HOST, PORT, *fem.mesh, x, "Final Position Solution with XAD AD"); + // + // std::println("Solver finished using XAD machine-precision gradients."); + // return 0; + +} \ No newline at end of file diff --git a/mapping_impl.ipp b/mapping_impl.ipp new file mode 100644 index 0000000..2fa0ad2 --- /dev/null +++ b/mapping_impl.ipp @@ -0,0 +1,216 @@ +FEM setup_fem(const std::string& filename, const bool verbose) { + FEM fem; + fem.mesh = std::make_unique(filename, 0, 0); + fem.mesh->EnsureNodes(); + int dim = fem.mesh->Dimension(); + fem.H1_fec = std::make_unique(2, dim); + fem.H1_fes = std::make_unique(fem.mesh.get(), fem.H1_fec.get()); + fem.Vec_H1_fes = std::make_unique(fem.mesh.get(), fem.H1_fec.get(), dim, mfem::Ordering::byNODES); + + fem.block_offsets.SetSize(3); + fem.block_offsets[0] = 0; + fem.block_offsets[1] = fem.H1_fes->GetTrueVSize(); + fem.block_offsets[2] = fem.H1_fes->GetTrueVSize() + fem.Vec_H1_fes->GetTrueVSize(); + fem.com.SetSize(dim); fem.com = 0.0; + fem.Q.SetSize(dim, dim); fem.Q = 0.0; + return fem; +} + +void view_mesh(const std::string& host, int port, const mfem::Mesh& mesh, const mfem::GridFunction& gf, const std::string& title) { + mfem::socketstream sol_sock(host.c_str(), port); + if (!sol_sock.is_open()) return; + sol_sock << "solution\n" << mesh << gf; + sol_sock << "window_title '" << title << "'\n" << std::flush; + +} + +double domain_integrate_grid_function(const FEM& fem, const mfem::GridFunction& gf) { + mfem::LinearForm lf(fem.H1_fes.get()); + mfem::GridFunctionCoefficient gf_c(&gf); + lf.AddDomainIntegrator(new mfem::DomainLFIntegrator(gf_c)); + lf.Assemble(); + return lf.Sum(); +} + +mfem::Vector get_com(const FEM& fem, const mfem::GridFunction &rho) { + const int dim = fem.mesh->Dimension(); + mfem::Vector com(dim); + com = 0.0; + double total_mass = 0.0; + + for (int i = 0; i < fem.H1_fes->GetNE(); ++i) { + mfem::ElementTransformation *trans = fem.H1_fes->GetElementTransformation(i); + const mfem::IntegrationRule &ir = mfem::IntRules.Get(trans->GetGeometryType(), fem.H1_fes->GetOrder(0) + trans->OrderW()); + + for (int j = 0; j < ir.GetNPoints(); ++j) { + const mfem::IntegrationPoint &ip = ir.IntPoint(j); + trans->SetIntPoint(&ip); + + double weight = trans->Weight() * ip.weight; + double rho_val = rho.GetValue(i, ip); + + mfem::Vector phys_point(dim); + trans->Transform(ip, phys_point); + + const double mass_term = rho_val * weight; + total_mass += mass_term; + + for (int d = 0; d < dim; ++d) { + com(d) += phys_point(d) * mass_term; + } + } + } + com /= total_mass; + return com; +} + +double centrifugal_potential(const mfem::Vector& x, double omega) { + const double s2 = std::pow(x(0), 2) + std::pow(x(1), 2); + return -0.5 * s2 * std::pow(omega, 2); +} + +double get_moment_of_inertia(const FEM& fem, const mfem::GridFunction& rho) { + auto s2_func = [](const mfem::Vector& x) { + return std::pow(x(0), 2) + std::pow(x(1), 2); + }; + + mfem::FunctionCoefficient s2_coeff(s2_func); + mfem::GridFunctionCoefficient rho_coeff(&rho); + + mfem::ProductCoefficient I_integrand ( rho_coeff, s2_coeff ); + mfem::LinearForm I_lf(fem.H1_fes.get()); + I_lf.AddDomainIntegrator(new mfem::DomainLFIntegrator(I_integrand)); + I_lf.Assemble(); + + return I_lf.Sum(); +} + +std::unique_ptr grav_potential(const FEM& fem, const Args &args, const mfem::GridFunction& rho) { + auto phi = std::make_unique(fem.H1_fes.get()); + + mfem::Array ess_bdr(fem.mesh->bdr_attributes.Max()); + ess_bdr = 1; + + mfem::GridFunctionCoefficient rho_coeff(&rho); + double total_mass = domain_integrate_grid_function(fem, rho); + + auto grav_potential = [&fem, &total_mass](const mfem::Vector& x) { + return l2_multipole_potential(fem, total_mass, x); + }; + + mfem::FunctionCoefficient phi_bdr_coeff(grav_potential); + mfem::Array ess_tdof_list; + fem.H1_fes->GetEssentialTrueDofs(ess_bdr, ess_tdof_list); + + auto laplacian = std::make_unique(fem.H1_fes.get()); + laplacian->AddDomainIntegrator(new mfem::DiffusionIntegrator()); + laplacian->Assemble(); + laplacian->Finalize(); + + mfem::ConstantCoefficient four_pi_G(-4.0 * M_PI * G); + mfem::ProductCoefficient rhs_coeff(rho_coeff, four_pi_G); + mfem::LinearForm b(fem.H1_fes.get()); + b.AddDomainIntegrator(new mfem::DomainLFIntegrator(rhs_coeff)); + b.Assemble(); + + mfem::OperatorPtr A; + mfem::Vector B, X; + laplacian->FormLinearSystem(ess_tdof_list, *phi, b, A, X, B); + + mfem::GSSmoother prec; + mfem::CGSolver cg; + cg.SetPreconditioner(prec); + cg.SetOperator(*A); + cg.SetRelTol(args.p.tol); + cg.SetMaxIter(args.p.max_iters); + cg.SetPrintLevel(0); + + cg.Mult(B, X); + laplacian->RecoverFEMSolution(X, b, *phi); + return phi; +} + +std::unique_ptr get_potential(const FEM &fem, const Args &args, const mfem::GridFunction &rho) { + auto phi = grav_potential(fem, args, rho); + mfem::GridFunctionCoefficient rho_coeff(&rho); + + if (args.r.enabled) { + auto rot = [&args](const mfem::Vector& x) { + return centrifugal_potential(x, args.r.omega); + }; + + mfem::FunctionCoefficient centrifugal_coeff(rot); + mfem::GridFunction centrifugal_gf(fem.H1_fes.get()); + centrifugal_gf.ProjectCoefficient(centrifugal_coeff); + (*phi) += centrifugal_gf; + } + + return phi; +} + +mfem::DenseMatrix compute_quadrupole_moment_tensor(const FEM& fem, const mfem::GridFunction& rho, const mfem::Vector& com) { + const int dim = fem.mesh->Dimension(); + mfem::DenseMatrix Q(dim, dim); + Q = 0.0; + + for (int i = 0; i < fem.H1_fes->GetNE(); ++i) { + mfem::ElementTransformation *trans = fem.mesh->GetElementTransformation(i); + const mfem::IntegrationRule &ir = mfem::IntRules.Get(trans->GetGeometryType(), 2 * fem.H1_fes->GetOrder(0) + trans->OrderW()); + for (int j = 0; j < ir.GetNPoints(); ++j) { + const mfem::IntegrationPoint &ip = ir.IntPoint(j); + trans->SetIntPoint(&ip); + + const double weight = trans->Weight() * ip.weight; + const double rho_val = rho.GetValue(i, ip); + + mfem::Vector phys_point(dim); + trans->Transform(ip, phys_point); + + mfem::Vector x_prime(dim); + double r_sq = 0.0; + + for (int d = 0; d < dim; ++d) { + x_prime(d) = phys_point(d) - com(d); + r_sq += x_prime(d) * x_prime(d); + } + + for (int m = 0; m < dim; ++m) { + for (int n = 0; n < dim; ++n) { + const double delta = (m == n) ? 1.0 : 0.0; + const double contrib = 3.0 * x_prime(m) * x_prime(n) - delta * r_sq; + Q(m, n) += rho_val * contrib * weight; + } + } + } + } + return Q; +} + +double l2_multipole_potential(const FEM &fem, const double total_mass, const mfem::Vector &x) { + const double r = x.Norml2(); + if (r < 1e-12) return 0.0; + + const int dim = fem.mesh->Dimension(); + + mfem::Vector n(x); + n /= r; + + double l2_mult_factor = 0.0; + for (int i = 0; i < dim; ++i) { + for (int j = 0; j < dim; ++j) { + l2_mult_factor += fem.Q(i, j) * n(i) * n(j); + } + } + + const double l2_contrib = (G / (2.0 * std::pow(r, 3))) * l2_mult_factor; + + const double l0_contrib = -G * total_mass / r; + + // l1 contribution is zero for a system centered on its COM + return l0_contrib + l2_contrib; +} + +void ConserveMass(const FEM& fem, mfem::GridFunction& rho, double target_mass) { + const double current_mass = domain_integrate_grid_function(fem, rho); + if (current_mass > 1e-15) rho *= (target_mass / current_mass); +} diff --git a/sweep_omega.sh b/sweep_omega.sh new file mode 100644 index 0000000..e69de29