diff --git a/.github/workflows/examples-smoke-run.yml b/.github/workflows/examples-smoke-run.yml index c3da5f5a..922a035b 100644 --- a/.github/workflows/examples-smoke-run.yml +++ b/.github/workflows/examples-smoke-run.yml @@ -44,7 +44,7 @@ jobs: -D ELASTICA_BUILD_BENCHMARKS=OFF \ -D ELASTICA_BUILD_DOCUMENTATION=OFF \ -D ELASTICA_BUILD_EXAMPLES=ON \ - -D ELASTICA_BUILD_PYTHON_BINDINGS=OFF + -D ELASTICA_BUILD_PYTHON_BINDINGS=ON - name: Build example executables run: | @@ -54,6 +54,7 @@ jobs: example_free_nest \ example_oscillatory_magnetic_field \ example_timoshenko \ + PyFishCaseStudy \ --parallel - name: Run Axial Stretching @@ -86,3 +87,14 @@ jobs: working-directory: build run: | ./bin/example_timoshenko --n_steps 10 + + - name: Install Fish Swimming Python dependencies + shell: bash + run: | + python -m pip install build/bin/python/ + python -m pip install examples/fish_swimming/ + + - name: Run Fish Swimming + shell: bash + run: | + python examples/fish_swimming/fish_swimming.py diff --git a/elastica/Utilities/AppSupport/CaseStudyInterface.hpp b/elastica/Utilities/AppSupport/CaseStudyInterface.hpp index 96aeb8a9..c41d4d0b 100644 --- a/elastica/Utilities/AppSupport/CaseStudyInterface.hpp +++ b/elastica/Utilities/AppSupport/CaseStudyInterface.hpp @@ -89,6 +89,15 @@ namespace elastica { } //************************************************************************** + //************************************************************************** + /*!\brief Change the time step to a given value and do one step. + */ + void step_with(::elastica::real_t dt) { + set_dt(dt); + self().step(); + } + //************************************************************************** + //@} //************************************************************************** diff --git a/elastica/Utilities/AppSupport/Python/CaseStudyPybind.hpp b/elastica/Utilities/AppSupport/Python/CaseStudyPybind.hpp index d62c5148..f31b0f75 100644 --- a/elastica/Utilities/AppSupport/Python/CaseStudyPybind.hpp +++ b/elastica/Utilities/AppSupport/Python/CaseStudyPybind.hpp @@ -118,6 +118,8 @@ namespace py_bindings { .def("run_for_n_iterations", &CaseStudy::run_for_n_iterations, py::arg("n_iterations"), py::pos_only(), "Runs the simulator for n additional iterations") + .def("step_with", &CaseStudy::step_with, py::arg("dt"), + "Change the time step to a given value and do one step.") .def("current_time", &CaseStudy::current_time, "Gets the current time of the simulation") .def("current_dt", &CaseStudy::current_dt, diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 070fd6ef..752bcbda 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -14,6 +14,7 @@ add_subdirectory(rod_plane_collision) add_subdirectory(nest) add_subdirectory(magnetism) +add_subdirectory(fish_swimming) # examples include(ElasticaAddExamples) diff --git a/examples/fish_swimming/.gitignore b/examples/fish_swimming/.gitignore new file mode 100644 index 00000000..b53b54b2 --- /dev/null +++ b/examples/fish_swimming/.gitignore @@ -0,0 +1,8 @@ +*.vtkhdf +*.pvsm +*.egg-info/ +build/ +dist/ +*.egg +*data*/ +*.so \ No newline at end of file diff --git a/examples/fish_swimming/CMakeLists.txt b/examples/fish_swimming/CMakeLists.txt new file mode 100644 index 00000000..e8127f70 --- /dev/null +++ b/examples/fish_swimming/CMakeLists.txt @@ -0,0 +1,42 @@ +add_subdirectory(FishModel) + +set(LIBRARY FishCaseStudy) + +add_library(${LIBRARY} STATIC ${LIBRARY}.cpp) +target_sources(${LIBRARY} INTERFACE + ${LIBRARY}.hpp + ) + +target_link_libraries(${LIBRARY} + PRIVATE + Block + SystemsCommon + CosseratRods + SystemsCustomization + RigidBody + States + Utilities + Simulator + Systems + Materials + ElasticaFlags + FishModel + ) + +# So that I can code up as rather than "CaseStudy" +target_include_directories(${LIBRARY} PRIVATE ${CMAKE_SOURCE_DIR}) + +include(ElasticaCaseStudyHelpers) + +add_elastica_case_study(example_elastic_fish + FROM ${LIBRARY} + AS elastica::standalone_executable + ) + +add_elastica_case_study(Py${LIBRARY} + FROM ${LIBRARY} + AS elastica::python_object + ) +# +elastica_python_link_libraries(Py${LIBRARY} PRIVATE Blaze) +elastica_python_add_dependencies(Py${LIBRARY} PyArrays PyCosseratRods PyFishModel) diff --git a/examples/fish_swimming/FishCaseStudy.cpp b/examples/fish_swimming/FishCaseStudy.cpp new file mode 100644 index 00000000..9a1f59d3 --- /dev/null +++ b/examples/fish_swimming/FishCaseStudy.cpp @@ -0,0 +1,133 @@ +#include "FishCaseStudy.hpp" + +#include +#include +#include + +void FishCaseStudy::setup(InputTypes const& input_parameters) { + ::elastica::MaterialID rod_material = elastica::create_material( + ::elastica::get(input_parameters), 0.24, + ::elastica::get(input_parameters), + ::elastica::get(input_parameters) / + ::elastica::real_t(3.0), + 200, 1e-10, 1e-10, 0.1, 0.1); + + const ::elastica::Vec3 origin{0.0, 0.0, 0.0}; + const ::elastica::real_t base_length{ + ::elastica::get(input_parameters)}; + const std::size_t n_elems = ::elastica::get(input_parameters); + const std::size_t n_fish = ::elastica::get(input_parameters); + std::size_t n_nodes = n_elems + 1UL; + + /* Initialize the fishes (actual + virtual) */ + using StraightInit = + ::elastica::cosserat_rod::StraightCosseratRodWithoutDampingInitializer; + for (std::size_t i = 0; i < n_fish; ++i) { + StraightInit straight_initializer{ + StraightInit::NElement{n_elems}, + StraightInit::Material{rod_material}, + StraightInit::Radius{base_length * 0.01}, // overwritten internally + StraightInit::Length{base_length}, + StraightInit::Origin{ + ::blaze::column(::elastica::get(input_parameters), i)}, + StraightInit::Direction{::elastica::Vec3{1.0, 0.0, 0.0}}, + StraightInit::Normal{::elastica::Vec3{0.0, 0.0, 1.0}}, + }; + + elastic_fish.emplace_back(simulator->template emplace_back( + straight_initializer.template initialize())); + virtual_fish.emplace_back(simulator->template emplace_back( + straight_initializer.template initialize())); + } + + /* Prescribe the virtual fishes' motion */ + for (auto& vf : virtual_fish) { + simulator->constrain(vf).at(0UL).with( + ::fish::constraint::make_carling_fish_constraint( + vf, + ::elastica::get(input_parameters), // period + 2 * M_PI / base_length, // wave number + 0.0, // phase shift + ::elastica::get(input_parameters) // ramp up time + )); + } + + /* Connect the virtual and actual fish */ + using Tag = ::elastica::tags::ShearStretchRigidityMatrix; + real_t const conn_k = ::elastica::get(elastic_fish[0].get())(2, 0); + real_t const y_idx = 1; + namespace econ = ::elastica::connections; + auto connection_generator = econ::make_connection( + econ::make_force([=](auto /*time*/, auto& suitor, + std::size_t /*rod_one_index*/, auto& leader, + auto /*rod_two_index*/) { + using ETag = ::elastica::tags::ExternalLoads; + ::blaze::row(::elastica::get(suitor), y_idx) += + conn_k * (::blaze::row(::elastica::position(leader), y_idx) - + ::blaze::row(::elastica::position(suitor), y_idx)); + }), + econ::make_torque([=](auto /*time*/, auto& suitor, + std::size_t /*rod_one_index*/, auto& leader, + auto /*rod_two_index*/) { + using TraitOp = + ::elastica::cosserat_rod::DefaultCosseratRodTraits::Operations; + using ITag = ::elastica::tags::InternalTorques; + using ETag = ::elastica::tags::ExternalTorques; + + // tau_1 + Q_1 * Q_2^T * m_2 = 0 + auto&& tmp1 = + ::elastica::get<::elastica::tags::_DummyElementVector>(leader); + auto&& tmp2 = + ::elastica::get<::elastica::tags::_DummyElementVector>(suitor); + TraitOp::batch_mattranspvec(tmp1, ::elastica::director(leader), + ::elastica::get(leader)); + TraitOp::batch_matvec(tmp2, ::elastica::director(suitor), tmp1); + ::elastica::get(suitor) -= tmp2; + })); + + using ::elastica::interface::at; + for (std::size_t i = 0; i < n_fish; ++i) { + simulator->connect(elastic_fish[i], at(0UL)) + .to(virtual_fish[i], at(0UL)) + .with(connection_generator()); + } + + /* Apply flow forces onto the elastic fish */ + namespace efo = ::elastica::forcing; + flow_forces = ::blaze::zero(3UL, n_fish * n_nodes); + flow_torques = ::blaze::zero(3UL, n_fish * n_elems); + for (std::size_t i = 0; i < n_fish; ++i) { + auto flow_force = + blaze::submatrix(flow_forces, 0UL, i * n_nodes, 3UL, n_nodes); + auto flow_torque = + blaze::submatrix(flow_torques, 0UL, i * n_elems, 3UL, n_elems); + simulator->force(elastic_fish[i]) + .at(0UL) + .with(efo::make_forcing( + efo::make_force([=](auto /*time*/, auto& rod_system, + std::size_t /*index*/) { + ::elastica::get<::elastica::tags::ExternalLoads>(rod_system) += + flow_force; + }), + efo::make_torque([=](auto /*time*/, auto& rod_system, + std::size_t /*index*/) { + ::elastica::get<::elastica::tags::ExternalTorques>(rod_system) += + flow_torque; + }))()); + } + + /* Finalize the simulator*/ + simulator->finalize(); + + /* Set the dt value */ + real_t dt = + ::elastica::get(input_parameters) * base_length / n_elems; + set_dt(dt); +} + +void FishCaseStudy::step() { + simulator->step(time_stepper, get_current_time(), get_dt(), + get_current_iteration()); +} + +void FishCaseStudy::run() {} diff --git a/examples/fish_swimming/FishCaseStudy.hpp b/examples/fish_swimming/FishCaseStudy.hpp new file mode 100644 index 00000000..686b3e89 --- /dev/null +++ b/examples/fish_swimming/FishCaseStudy.hpp @@ -0,0 +1,169 @@ +#pragma once +// +#include +// +#include +#include +#include +#include +#include +#include +#include +#include +// +#include +// +#include // transform +#include // make_tuple +#include + +using String = const char* const; +namespace ec = elastica::configuration; +using RodType = ::fish::Fish; +using ::elastica::real_t; +using ::elastica::Vec3; + +struct FishConfiguration { + static inline auto make_simulator_configuration() { + auto system_config = ec::make_systems_configuration>( + ec::make_blocking_policy(ec::RestrictSizeAcrossBlockTypes{50UL}), + ec::make_iteration_policy(ec::iterate_across_system_types_in(ec::seq), + ec::iterate_over_systems_in(ec::seq))); + + return ec::make_simulator_configuration( + ec::DefaultDomainConfiguration{}, system_config, + ec::DefaultParallelConfiguration{}, + ec::DefaultConstraintsConfiguration{}, + ec::DefaultConnectionsConfiguration{}, + ec::DefaultForcingConfiguration{}); + } + static inline auto make_simulator() { + return ::elastica::make_simulator<::elastica::modules::Constraints, + ::elastica::modules::Connections, + ::elastica::modules::Forcing>( + make_simulator_configuration()); + } + + static inline auto make_time_stepper() { + using TimeStepper = ::elastica::timestepper::PositionVerlet; + return TimeStepper{}; + } +}; + +struct FishCaseStudy : public ::elastica::CaseStudyInterface { + using Configuration = FishConfiguration; + using This = FishCaseStudy; + + static constexpr String help = {"Fish case study"}; + + struct num_fish { + using type = std::size_t; + static constexpr String help = {"Number of fish"}; + static constexpr type default_value() { return 1UL; } + }; + + struct origins { + using type = ::blaze::DynamicMatrix; // 3 x num_fish + static constexpr String help = {"Origins of the fish"}; + static type default_value() { return {1UL, 3UL, 0.0}; } + }; + + struct n_elements { + using type = std::size_t; + static constexpr String help = {"Number of elements in the rod"}; + static constexpr type default_value() { return 50UL; } + }; + + struct density { + using type = real_t; + static constexpr String help = {"Density of the rod"}; + static constexpr type default_value() { return 66.0; } + }; + + struct youngs_modulus { + using type = real_t; + static constexpr String help = {"Youngs modulus of the rod"}; + static constexpr type default_value() { return 1.5e6; } + }; + + struct damping_rate { + using type = real_t; + static constexpr String help = {"Rate of damping in the rod"}; + static constexpr type default_value() { return 0.02; } + }; + + struct length { + using type = real_t; + static constexpr String help = {"Length of the rod"}; + static constexpr type default_value() { return 1.0; } + }; + + struct dt_scale { + using type = real_t; + static constexpr String help = {"Scale for time steps"}; + static constexpr type default_value() { return 5e-4; } + }; + + struct period { + using type = real_t; + static constexpr String help = {"Period of the rod's motion"}; + static constexpr type default_value() { return 1.0; } + }; + + using Parameters = + tmpl::list; + + /* Required functions */ + void setup(InputTypes const& input_parameters); + void step(); + void run(); + + /* Python communication function */ + template + inline auto communicate_buffer( + std::vector> const& ref_generators) { + using CT = + decltype(std::declval const&>().get()); + using RT = std::vector; + RT return_list; + return_list.reserve(ref_generators.size()); + std::transform(std::begin(ref_generators), std::end(ref_generators), + std::back_inserter(return_list), + [](blocks::BlockSliceGenerator const& rod_ref) { + return rod_ref.get(); + }); + return return_list; + } + + inline auto communicate() { + auto flow_force = ::blaze::submatrix( + flow_forces, 0UL, 0UL, flow_forces.rows(), flow_forces.columns()); + auto flow_torque = ::blaze::submatrix( + flow_torques, 0UL, 0UL, flow_torques.rows(), flow_torques.columns()); + return std::make_tuple(communicate_buffer(elastic_fish), + communicate_buffer(virtual_fish), flow_force, + flow_torque); + } + + FishCaseStudy() = default; + FishCaseStudy(FishCaseStudy const&) = default; + FishCaseStudy(FishCaseStudy&&) noexcept = default; + FishCaseStudy& operator=(FishCaseStudy const&) = default; + FishCaseStudy& operator=(FishCaseStudy&&) noexcept = default; + + decltype(Configuration::make_time_stepper()) time_stepper = + Configuration::make_time_stepper(); + decltype(Configuration::make_simulator()) simulator = + Configuration::make_simulator(); + + // Class members + using RodRef = decltype(simulator->get_invalid_ref()); + std::vector elastic_fish; + std::vector virtual_fish; + + ::blaze::DynamicMatrix flow_forces; + ::blaze::DynamicMatrix flow_torques; +}; + +ELASTICA_REGISTER_CASE_STUDY(FishCaseStudy); diff --git a/examples/fish_swimming/FishModel/Aliases.hpp b/examples/fish_swimming/FishModel/Aliases.hpp new file mode 100644 index 00000000..6fb44ee8 --- /dev/null +++ b/examples/fish_swimming/FishModel/Aliases.hpp @@ -0,0 +1,3 @@ +#pragma once + +#include diff --git a/examples/fish_swimming/FishModel/Aliases/CMakeLists.txt b/examples/fish_swimming/FishModel/Aliases/CMakeLists.txt new file mode 100644 index 00000000..9f72bc02 --- /dev/null +++ b/examples/fish_swimming/FishModel/Aliases/CMakeLists.txt @@ -0,0 +1,18 @@ +set(LIBRARY FishModelAliases) + +add_elastica_library(${LIBRARY} INTERFACE) + +set(LIBRARY_SOURCES + FishCosseratRod.hpp +) +elastica_target_headers( + ${LIBRARY} + INCLUDE_DIRECTORY ${CMAKE_SOURCE_DIR}/examples + ${LIBRARY_SOURCES} +) + +elastica_target_sources( + ${LIBRARY} + INTERFACE + ${LIBRARY_SOURCES} +) \ No newline at end of file diff --git a/examples/fish_swimming/FishModel/Aliases/FishCosseratRod.hpp b/examples/fish_swimming/FishModel/Aliases/FishCosseratRod.hpp new file mode 100644 index 00000000..a58792c6 --- /dev/null +++ b/examples/fish_swimming/FishModel/Aliases/FishCosseratRod.hpp @@ -0,0 +1,49 @@ +#pragma once + +// common +#include +// blocks +#include +// cosserat rod components +#include +#include + +namespace fish { + + namespace details { + + template + using FishComponents = ::elastica::SymplecticPolicy< + T, B, ::fish::component::WithFishGeometry, + ::elastica::cosserat_rod::component:: + WithDiagonalLinearHyperElasticModel, + ::elastica::cosserat_rod::component::WithRodKinematics, + ::elastica::cosserat_rod::component::CosseratRodNameAdapted>; + + using Fish = ::elastica::cosserat_rod::CosseratRodPlugin< + ::elastica::cosserat_rod::DefaultCosseratRodTraits, ::blocks::Block, + FishComponents>; + using FishSlice = ::elastica::cosserat_rod::CosseratRodPlugin< + ::elastica::cosserat_rod::DefaultCosseratRodTraits, + ::blocks::BlockSlice, FishComponents>; + using FishConstSlice = ::elastica::cosserat_rod::CosseratRodPlugin< + ::elastica::cosserat_rod::DefaultCosseratRodTraits, + ::blocks::ConstBlockSlice, FishComponents>; + using FishView = ::elastica::cosserat_rod::CosseratRodPlugin< + ::elastica::cosserat_rod::DefaultCosseratRodTraits, ::blocks::BlockView, + FishComponents>; + using FishConstView = ::elastica::cosserat_rod::CosseratRodPlugin< + ::elastica::cosserat_rod::DefaultCosseratRodTraits, + ::blocks::ConstBlockView, FishComponents>; + + using FishBlock = ::blocks::Block; + using FishBlockSlice = ::blocks::BlockSlice; + using FishBlockConstSlice = ::blocks::ConstBlockSlice; + using FishBlockView = ::blocks::BlockView; + using FishBlockConstView = ::blocks::ConstBlockView; + + } // namespace details + + using Fish = details::Fish; + +} // namespace fish diff --git a/examples/fish_swimming/FishModel/CMakeLists.txt b/examples/fish_swimming/FishModel/CMakeLists.txt new file mode 100644 index 00000000..6a8f4b10 --- /dev/null +++ b/examples/fish_swimming/FishModel/CMakeLists.txt @@ -0,0 +1,41 @@ +set(LIBRARY FishModel) + +add_elastica_library(${LIBRARY} INTERFACE) + +set(LIBRARY_SOURCES + Tags.hpp + Aliases.hpp + Components.hpp + Constraints.hpp + Dissipations.hpp + FishModel.hpp +) + +elastica_target_headers( + ${LIBRARY} + INCLUDE_DIRECTORY ${CMAKE_SOURCE_DIR}/examples + ${LIBRARY_SOURCES} +) + +elastica_target_sources( + ${LIBRARY} + INTERFACE + ${LIBRARY_SOURCES} +) + +add_subdirectory(Aliases) +add_library(${LIBRARY}::Aliases ALIAS FishModelAliases) + +add_subdirectory(Components) +add_library(${LIBRARY}::Components ALIAS FishModelComponents) + +add_subdirectory(Constraints) +add_library(${LIBRARY}::Components ALIAS FishModelConstraints) + +# add_subdirectory(Connections) +# add_library(${LIBRARY}::Components ALIAS FishModelConnections) + +add_subdirectory(Dissipations) +add_library(${LIBRARY}::Components ALIAS FishModelDissipations) + +add_subdirectory(Python) diff --git a/examples/fish_swimming/FishModel/Components.hpp b/examples/fish_swimming/FishModel/Components.hpp new file mode 100644 index 00000000..1e985d61 --- /dev/null +++ b/examples/fish_swimming/FishModel/Components.hpp @@ -0,0 +1,3 @@ +#pragma once + +#include diff --git a/examples/fish_swimming/FishModel/Components/CMakeLists.txt b/examples/fish_swimming/FishModel/Components/CMakeLists.txt new file mode 100644 index 00000000..48398baa --- /dev/null +++ b/examples/fish_swimming/FishModel/Components/CMakeLists.txt @@ -0,0 +1,18 @@ +set(LIBRARY FishModelComponents) + +add_elastica_library(${LIBRARY} INTERFACE) + +set(LIBRARY_SOURCES + WithFishGeometry.hpp +) +elastica_target_headers( + ${LIBRARY} + INCLUDE_DIRECTORY ${CMAKE_SOURCE_DIR}/examples + ${LIBRARY_SOURCES} +) + +elastica_target_sources( + ${LIBRARY} + INTERFACE + ${LIBRARY_SOURCES} +) \ No newline at end of file diff --git a/examples/fish_swimming/FishModel/Components/WithFishGeometry.hpp b/examples/fish_swimming/FishModel/Components/WithFishGeometry.hpp new file mode 100644 index 00000000..ef462119 --- /dev/null +++ b/examples/fish_swimming/FishModel/Components/WithFishGeometry.hpp @@ -0,0 +1,231 @@ +#pragma once + +#include +#include +#include +#include +#include +// +#include +#include +// +#include + +namespace fish { + + namespace component { + + namespace details { + + template + struct FishGeometryVariables { + private: + using Traits = CRT; + + protected: + struct SemiMajorAxis : public Traits::template CosseratRodVariable< + ::fish::tags::SemiMajorAxis, + typename Traits::DataType::Scalar, + typename Traits::Place::OnElement> {}; + struct SemiMinorAxis : public Traits::template CosseratRodVariable< + ::fish::tags::SemiMinorAxis, + typename Traits::DataType::Scalar, + typename Traits::Place::OnElement> {}; + + using ComputedVariables = tmpl::list; + using Variables = ComputedVariables; + }; + + } // namespace details + + template + class WithFishGeometry + : public ::elastica::CRTPHelper, + public ::elastica::cosserat_rod::component::detail:: + CosseratRodCrossSectionInterface< + CRT, ComputationalBlock, + ::elastica::cosserat_rod::component::detail:: + CircleCrossSectionOperations>, + public details::FishGeometryVariables, + public ::tt::ConformsTo< + ::elastica::cosserat_rod::component::protocols::Geometry1D>, + public ::elastica::GeometryComponent< + WithFishGeometry> { + private: + using Traits = CRT; + using CRTP = ::elastica::CRTPHelper; + using CrossSection = ::elastica::cosserat_rod::component::detail:: + CircleCrossSectionOperations; + using This = WithFishGeometry; + using Parent = ::elastica::cosserat_rod::component::detail:: + CosseratRodCrossSectionInterface; + using VariableDefinitions = + typename details::FishGeometryVariables; + using real_t = typename Traits::real_type; + using index_t = typename Traits::index_type; + + // sanity check cross section classes + static_assert( + std::is_same_v, + "Cross section assertion failure, contact the developers!"); + + protected: + using typename Parent::DirectionType; + using typename Parent::Frame; + using typename Parent::InitializedVariables; + using typename VariableDefinitions::SemiMajorAxis; + using typename VariableDefinitions::SemiMinorAxis; + using ComputedVariables = + tmpl::append; + using Variables = tmpl::append; + + template + static void initialize( + WithFishGeometry& this_component, + CosseratInitializer&& initializer) { + // 1. Initialize parent + Parent::initialize(this_component, + std::forward(initializer)); + + // 2. Initialize required variables + ::elastica::cosserat_rod::initialize_component( + this_component.self(), + std::forward(initializer)); + + // 3. Compute other variables + { + auto&& reference_lengths( + blocks::get<::elastica::tags::ReferenceElementLength>( + this_component.self())); + auto&& semi_major_axis( + blocks::get<::fish::tags::SemiMajorAxis>(this_component.self())); + auto&& semi_minor_axis( + blocks::get<::fish::tags::SemiMinorAxis>(this_component.self())); + + const real_t base_length = ::blaze::sum(reference_lengths); + const index_t n_elements = reference_lengths.size(); + + const real_t s_b{0.04 * base_length}; + const real_t s_t{0.95 * base_length}; + const real_t w_h{0.04 * base_length}; + const real_t w_t{0.01 * base_length}; + const real_t a{0.51 * base_length}; + const real_t b{0.08 * base_length}; + + const real_t dx{1.0 / n_elements}; + + for (index_t i = 0; i < n_elements; ++i) { + real_t s = (static_cast(i) + 0.5) * dx; + if (s >= 0 && s <= s_b) { + semi_minor_axis[i] = std::sqrt(2 * w_h * s - s * s); + } else if (s_b < s && s <= s_t) { + semi_minor_axis[i] = + w_h - (w_h - w_t) * std::pow((s - s_b) / (s_t - s_b), 2); + } else if (s_t < s && s <= base_length) { + semi_minor_axis[i] = + w_t * (base_length - s) / (base_length - s_t); + } else { + semi_minor_axis[i] = 0.0; + } + + semi_major_axis[i] = b * std::sqrt(1 - std::pow((s - a) / a, 2)); + } + + auto&& element_dimensions( + blocks::get<::elastica::tags::ElementDimension>( + this_component.self())); + auto&& volumes(blocks::get<::elastica::tags::ElementVolume>( + this_component.self())); + volumes = + M_PI * semi_major_axis * semi_minor_axis * reference_lengths; + element_dimensions = + ::elastica::sqrt(volumes / reference_lengths / M_PI); + } + } + + public: + using CRTP::self; + + // Get methods: SemiMajorAxis + inline constexpr decltype(auto) get_semi_major_axis() & noexcept { + return blocks::get<::fish::tags::SemiMajorAxis>(self()); + } + inline constexpr decltype(auto) get_semi_major_axis() const& noexcept { + return blocks::get<::fish::tags::SemiMajorAxis>(self()); + } + inline constexpr decltype(auto) get_semi_major_axis( + index_t idx) & noexcept { + return SemiMajorAxis::slice(get_semi_major_axis(), idx); + } + inline constexpr decltype(auto) get_semi_major_axis( + index_t idx) const& noexcept { + return SemiMajorAxis::slice(get_semi_major_axis(), idx); + } + + // Get methods: SemiMinorAxis + inline constexpr decltype(auto) get_semi_minor_axis() & noexcept { + return blocks::get<::fish::tags::SemiMinorAxis>(self()); + } + inline constexpr decltype(auto) get_semi_minor_axis() const& noexcept { + return blocks::get<::fish::tags::SemiMinorAxis>(self()); + } + inline constexpr decltype(auto) get_semi_minor_axis( + index_t idx) & noexcept { + return SemiMinorAxis::slice(get_semi_minor_axis(), idx); + } + inline constexpr decltype(auto) get_semi_minor_axis( + index_t idx) const& noexcept { + return SemiMinorAxis::slice(get_semi_minor_axis(), idx); + } + + // Get methods: Area + inline constexpr auto get_area() const& COSSERATROD_LIB_NOEXCEPT { + return M_PI * get_semi_major_axis() * get_semi_minor_axis(); + } + inline constexpr auto get_area( + index_t idx) const& COSSERATROD_LIB_NOEXCEPT { + return M_PI * get_semi_major_axis(idx) * get_semi_minor_axis(idx); + } + + // Get methods: Second Moment of Area + template = nullptr> + inline constexpr decltype(auto) get_second_moment_of_area( + index_t idx) const COSSERATROD_LIB_NOEXCEPT { + // 0.25 * pi * w^3 * h + return 0.25 * M_PI * get_semi_major_axis(idx) * + std::pow(get_semi_minor_axis(idx), 3); + } + template = nullptr> + inline constexpr decltype(auto) get_second_moment_of_area( + index_t idx) const COSSERATROD_LIB_NOEXCEPT { + // 0.25 * pi * w * h^3 + return 0.25 * M_PI * get_semi_minor_axis(idx) * + std::pow(get_semi_major_axis(idx), 3); + } + template = nullptr> + inline constexpr decltype(auto) get_second_moment_of_area( + index_t idx) const COSSERATROD_LIB_NOEXCEPT { + // 0.25 * pi * w * h * (w^2 + h^2) + return This::get_second_moment_of_area(idx) + + This::get_second_moment_of_area(idx); + } + inline constexpr auto get_second_moment_of_area(index_t idx) const + COSSERATROD_LIB_NOEXCEPT { + return ::elastica::Vec3{ + real_t(This::get_second_moment_of_area(idx)), + real_t(This::get_second_moment_of_area(idx)), + real_t(This::get_second_moment_of_area(idx))}; + } + + static inline constexpr auto shape_factor() COSSERATROD_LIB_NOEXCEPT + -> real_t { + return 27.0 / 28.0; + } + }; + + } // namespace component + +} // namespace fish diff --git a/examples/fish_swimming/FishModel/Constraints.hpp b/examples/fish_swimming/FishModel/Constraints.hpp new file mode 100644 index 00000000..15cb6027 --- /dev/null +++ b/examples/fish_swimming/FishModel/Constraints.hpp @@ -0,0 +1,3 @@ +#pragma once + +#include \ No newline at end of file diff --git a/examples/fish_swimming/FishModel/Constraints/CMakeLists.txt b/examples/fish_swimming/FishModel/Constraints/CMakeLists.txt new file mode 100644 index 00000000..d470d2de --- /dev/null +++ b/examples/fish_swimming/FishModel/Constraints/CMakeLists.txt @@ -0,0 +1,18 @@ +set(LIBRARY FishModelConstraints) + +add_elastica_library(${LIBRARY} INTERFACE) + +set(LIBRARY_SOURCES + CarlingFishConstraint.hpp +) +elastica_target_headers( + ${LIBRARY} + INCLUDE_DIRECTORY ${CMAKE_SOURCE_DIR}/examples + ${LIBRARY_SOURCES} +) + +elastica_target_sources( + ${LIBRARY} + INTERFACE + ${LIBRARY_SOURCES} +) \ No newline at end of file diff --git a/examples/fish_swimming/FishModel/Constraints/CarlingFishConstraint.hpp b/examples/fish_swimming/FishModel/Constraints/CarlingFishConstraint.hpp new file mode 100644 index 00000000..a472f25f --- /dev/null +++ b/examples/fish_swimming/FishModel/Constraints/CarlingFishConstraint.hpp @@ -0,0 +1,120 @@ +#pragma once + +#include +#include +// +#include +#include +#include +// +#include +#include +#include +// +#include + +namespace fish { + namespace constraint { + namespace details { + template class ComputationalBlock, + template class... Components> + decltype(auto) make_carling_fish_constraint( + ::elastica::cosserat_rod::CosseratRodPlugin< + Traits, ComputationalBlock, Components...> const& rod_like, + ::elastica::real_t const period, ::elastica::real_t const wave_number, + ::elastica::real_t const phase_shift, + ::elastica::real_t const ramp_up_time) { + namespace eco = ::elastica::constraints; + using ::elastica::real_t; + using ::elastica::Rot3; + using ::elastica::Vec3; + + auto const& rod = rod_like.self(); + + auto const n_elems{std::size(rod)}; + real_t const length{::blaze::sum( + ::elastica::get<::elastica::tags::ReferenceElementLength>(rod))}; + real_t const total_mass{ + ::blaze::sum(::elastica::get<::elastica::tags::Mass>(rod))}; + real_t const angular_frequency{2 * M_PI / period}; + Vec3 const start_position{::elastica::position(rod, 0UL)}; + + std::size_t const x_idx = 0UL; + std::size_t const y_idx = 1UL; + + return eco::make_constraint( + eco::make_value_constraint([=](auto time, auto& rod_system, + std::size_t /*index*/) { + auto&& position{::elastica::position(rod_system)}; + auto&& director{::elastica::director(rod_system)}; + auto&& mass{::blocks::get<::elastica::tags::Mass>(rod_system)}; + Vec3 com = position * mass / total_mass; + + real_t ramp_factor = + time <= ramp_up_time + ? 0.5 * (1 + std::sin(M_PI * (time / ramp_up_time - 0.5))) + : 1.0; + + // reset position for update + position *= 0.0; + + real_t const ds0 = 1.0 / static_cast(n_elems); + real_t const ds2 = std::pow(length * ds0, 2); + real_t s = 0.0; + + for (std::size_t i = 0UL; i < n_elems + 1UL; ++i) { + // Gazzola et al. JCP, 2011 Eqn 45 + position(y_idx, i) = + 0.125 * ramp_factor * length * (0.03125 + s) / 1.03125 * + std::sin(wave_number * s - angular_frequency * time + + phase_shift); + // (delta x)^2 + (delta y)^2 = (ds)^2 + if (i != 0) { + position(x_idx, i) = + position(x_idx, i - 1) + + std::sqrt(ds2 - std::pow(position(y_idx, i) - + position(y_idx, i - 1), + 2)); + } + + s += ds0; + } + + // Adjust position + for (std::size_t i = 0UL; i < n_elems + 1UL; ++i) { + ::blaze::column(position, i) += start_position; + } + real_t y_com = ::blaze::dot(mass, ::blaze::row(position, y_idx)) / + total_mass; + ::blaze::row(position, y_idx) += com[y_idx] - y_com; + + // Update directors + auto&& position_diff( + Traits::Operations::difference_kernel(position)); + for (std::size_t i = 0UL; i < n_elems; ++i) { + ::blaze::columnslice(director, i) = + ::elastica::make_orthogonal_bases_from_normal_and_tangent( + Vec3{0, 0, 1}, // shouldn't hardcode this + ::blaze::normalize(::blaze::column(position_diff, i))); + } + }), + eco::make_rate_constraint( + [](auto /*time*/, auto& rod_system, std::size_t /*index*/) { + ::elastica::velocity(rod_system) *= 0.0; + ::elastica::angular_velocity(rod_system) *= 0.0; + }))(); + } + } // namespace details + + /* Make periodic fish motion via constraint. Modeled after Carling et al., + * 1998. */ + template + decltype(auto) make_carling_fish_constraint( + ::blocks::BlockGenerator const& generator, + Args&&... args) { + return details::make_carling_fish_constraint( + blocks::generate_from(generator), std::forward(args)...); + } + } // namespace constraint +} // namespace fish \ No newline at end of file diff --git a/examples/fish_swimming/FishModel/Dissipations.hpp b/examples/fish_swimming/FishModel/Dissipations.hpp new file mode 100644 index 00000000..8aa66c74 --- /dev/null +++ b/examples/fish_swimming/FishModel/Dissipations.hpp @@ -0,0 +1,3 @@ +#pragma once + +#include \ No newline at end of file diff --git a/examples/fish_swimming/FishModel/Dissipations/AnalyticalLinearDamper.hpp b/examples/fish_swimming/FishModel/Dissipations/AnalyticalLinearDamper.hpp new file mode 100644 index 00000000..a1cb0d20 --- /dev/null +++ b/examples/fish_swimming/FishModel/Dissipations/AnalyticalLinearDamper.hpp @@ -0,0 +1,193 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace fish { + + namespace detail { + + /* Physical version */ + template class ComputationalBlock, + template class... Components> + decltype(auto) make_damping_constraint( + ::elastica::cosserat_rod::CosseratRodPlugin< + Traits, ComputationalBlock, Components...> const& rod_like, + ::elastica::real_t translational_damping_coeff, + ::elastica::real_t rotational_damping_coeff, ::elastica::real_t dt) { + namespace eco = ::elastica::constraints; + using ::elastica::real_t; + auto const& rod = rod_like.self(); + + auto&& nodal_mass{::elastica::get<::elastica::tags::Mass>(rod)}; + auto&& inv_moi{ + ::elastica::get<::elastica::tags::InvMassSecondMomentOfInertia>(rod)}; + + constexpr std::size_t dim{3UL}; + const auto n_elem{std::size(rod)}; + const auto n_node{n_elem + 1UL}; + + // Compute exponential damping constants + // Rotational damping constants is still missing dilatation term, + // which is dynamically computed in the constraint lambdas + blaze::DynamicVector translational_damping_constant(n_node); + blaze::DynamicMatrix rotational_damping_constant(dim, n_elem); + + for (std::size_t j = 0UL; j < n_node; ++j) { + translational_damping_constant[j] = + std::exp(-translational_damping_coeff * dt); + } + + for (std::size_t i = 0UL; i < dim; ++i) { + for (std::size_t j = 0UL; j < n_elem; ++j) { + real_t element_mass = 0.5 * (nodal_mass[j] + nodal_mass[j + 1]); + if (j == 0) + element_mass += 0.5 * nodal_mass[j]; + if (j == (n_elem - 1)) + element_mass += 0.5 * nodal_mass[j + 1]; + + rotational_damping_constant(i, j) = std::exp( + -rotational_damping_coeff * element_mass * inv_moi(i, j) * dt); + } + } + + return eco::make_constraint( + eco::make_value_constraint([](auto /*time*/, auto& /*rod_system*/, + std::size_t /*index*/) {}), + eco::make_rate_constraint([=](auto /*time*/, auto& rod_system, + std::size_t) { + auto&& vel = ::elastica::velocity(rod_system); + auto&& omega = ::elastica::angular_velocity(rod_system); + auto&& dil = + ::blocks::get(rod_system); + + for (std::size_t i = 0UL; i < dim; ++i) { + for (std::size_t j = 0UL; j < n_elem + 1UL; ++j) { + vel(i, j) *= translational_damping_constant[j]; + } + + for (auto j = 0UL; j < n_elem; ++j) { + omega(i, j) *= + std::pow(rotational_damping_constant(i, j), dil[j]); + } + } + }))(); + } + + /* Simplified version */ + template class ComputationalBlock, + template class... Components> + decltype(auto) make_damping_constraint( + ::elastica::cosserat_rod::CosseratRodPlugin< + Traits, ComputationalBlock, Components...> const& rod_like, + ::elastica::real_t damping_coeff, ::elastica::real_t dt) { + namespace eco = ::elastica::constraints; + using ::elastica::real_t; + auto const& rod = rod_like.self(); + + const real_t damping_constant{std::exp(-damping_coeff * dt)}; + + return eco::make_constraint( + eco::make_value_constraint([](auto /*time*/, auto& /*rod_system*/, + std::size_t /*index*/) {}), + eco::make_rate_constraint( + [damping_constant](auto /*time*/, auto& rod_system, std::size_t) { + ::elastica::velocity(rod_system) *= damping_constant; + ::elastica::angular_velocity(rod_system) *= damping_constant; + }))(); + } + + /* Original version */ + template class ComputationalBlock, + template class... Components> + decltype(auto) make_damping_constraint( + ::elastica::cosserat_rod::CosseratRodPlugin< + Traits, ComputationalBlock, Components...> const& rod_like, + ::elastica::real_t damping_coeff_dt) { + namespace eco = ::elastica::constraints; + using ::elastica::real_t; + auto const& rod = rod_like.self(); + + const real_t translational_mode_damping_coeff = + std::exp(-damping_coeff_dt); + constexpr std::size_t dim{3UL}; + const auto n_elem(std::size(rod)); + + auto&& mass = ::elastica::get(rod); + + blaze::DynamicMatrix rotational_mode_damping_coeff(dim, n_elem); + + for (std::size_t column_idx = 0; column_idx < n_elem; column_idx++) { + real_t elemental_mass = 0.5 * (mass[column_idx] + mass[column_idx + 1]); + if (column_idx == 0) + elemental_mass += 0.5 * mass[column_idx]; + if (column_idx == (n_elem - 1)) + elemental_mass += 0.5 * mass[column_idx + 1]; + for (std::size_t row_idx = 0; row_idx < dim; row_idx++) { + rotational_mode_damping_coeff(row_idx, column_idx) = std::exp( + -damping_coeff_dt * elemental_mass * + ::elastica::get<::elastica::tags::InvMassSecondMomentOfInertia>( + rod)(row_idx, column_idx)); + } + } + + return eco::make_constraint( + eco::make_value_constraint([](auto /*time*/, auto& /*rod_system*/, + std::size_t /*index*/) {}), + eco::make_rate_constraint([translational_mode_damping_coeff, + rotational_mode_damping_coeff]( + auto /*time*/, auto& rod_system, + std::size_t) { + ::elastica::velocity(rod_system) *= + translational_mode_damping_coeff; + // Easier to encode as powers + auto const n_elem(std::size(rod_system)); + auto&& omega = ::elastica::angular_velocity(rod_system); + auto&& dil = + ::blocks::get(rod_system); + + for (auto idx_elem = 0UL; idx_elem < n_elem; ++idx_elem) { + omega(0UL, idx_elem) *= std::pow( + rotational_mode_damping_coeff(0UL, idx_elem), dil[idx_elem]); + omega(1UL, idx_elem) *= std::pow( + rotational_mode_damping_coeff(1UL, idx_elem), dil[idx_elem]); + omega(2UL, idx_elem) *= std::pow( + rotational_mode_damping_coeff(2UL, idx_elem), dil[idx_elem]); + } + }))(); + } + + } // namespace detail + + /* Usage + // simplified version + simulator->constrain(block_slice_or_view) + .at(0UL) // index is arbitrary + .with(make_damping_constraint(block_slice_or_view, nu, dt)); + + // physical version + simulator->constrain(block_slice_or_view) + .at(0UL) // index is arbitrary + .with(make_damping_constraint(block_slice_or_view, nu_t, nu_r, dt)); + + // Original version + simulator->constrain(block_slice_or_view) + .at(0UL) // index is arbitrary + .with(make_damping_constraint(block_slice_or_view, nu*dt)); + */ + template + decltype(auto) make_damping_constraint( + ::blocks::BlockGenerator const& generator, + Args&&... args) { + return detail::make_damping_constraint(blocks::generate_from(generator), + std::forward(args)...); + } + +} // namespace fish diff --git a/examples/fish_swimming/FishModel/Dissipations/CMakeLists.txt b/examples/fish_swimming/FishModel/Dissipations/CMakeLists.txt new file mode 100644 index 00000000..fbd7c611 --- /dev/null +++ b/examples/fish_swimming/FishModel/Dissipations/CMakeLists.txt @@ -0,0 +1,18 @@ +set(LIBRARY FishModelDissipations) + +add_elastica_library(${LIBRARY} INTERFACE) + +set(LIBRARY_SOURCES + AnalyticalLinearDamper.hpp +) +elastica_target_headers( + ${LIBRARY} + INCLUDE_DIRECTORY ${CMAKE_SOURCE_DIR}/examples + ${LIBRARY_SOURCES} +) + +elastica_target_sources( + ${LIBRARY} + INTERFACE + ${LIBRARY_SOURCES} +) \ No newline at end of file diff --git a/examples/fish_swimming/FishModel/FishModel.hpp b/examples/fish_swimming/FishModel/FishModel.hpp new file mode 100644 index 00000000..09b84471 --- /dev/null +++ b/examples/fish_swimming/FishModel/FishModel.hpp @@ -0,0 +1,8 @@ +#pragma once + +#include +#include +#include +#include +#include +#include diff --git a/examples/fish_swimming/FishModel/Python/BindFishModel.cpp b/examples/fish_swimming/FishModel/Python/BindFishModel.cpp new file mode 100644 index 00000000..00dfea30 --- /dev/null +++ b/examples/fish_swimming/FishModel/Python/BindFishModel.cpp @@ -0,0 +1,85 @@ +#include +#include +#include +#include +// +#include +// +#include +// +#include // slice() +#include +#include +#include +#include +#include +// +#include +#include +#include +// +#include + +namespace py = pybind11; + +namespace py_bindings { + void bind_fish_model(py::module& m) { // NOLINT + using UnwantedVariableTags = + tmpl::list<::elastica::tags::_DummyElementVector, + ::elastica::tags::_DummyElementVector2, + ::elastica::tags::_DummyVoronoiVector>; + + { + using type = ::fish::details::FishBlockSlice; + + // Wrapper for basic type operations + auto py_cosserat_rod = + py::class_(m, "Fish") + .def("__len__", [](const type& t) { return std::size(t); }) + .def( + "__str__", + +[](const type& /*meta*/) { return std::string("Fish"); }); + bind_member_access(py_cosserat_rod); + bind_variable_locator(py_cosserat_rod); + bind_pyelastica_interface(py_cosserat_rod); + // bind_energy_methods(py_cosserat_rod); + } + { + using type = ::fish::details::FishBlockView; + + // Wrapper for basic type operations + auto py_cosserat_rod = + py::class_(m, "FishView") + .def("__len__", [](const type& t) { return std::size(t); }) + .def( + "__str__", + +[](const type& /*meta*/) { return std::string("FishView"); }) + .def( + "__getitem__", + +[](type& t, std::size_t index) { + variable_bounds_check(::blocks::n_units(t), index); + return ::blocks::slice(t, index); + }) + .def( + "__getitem__", +[](type& t, const py::slice slice) { + auto v = check_slice<0UL>(::blocks::n_units(t), + std::move(slice)); + return ::blocks::slice(t, v.start, v.start + v.slicelength); + }); + bind_member_access(py_cosserat_rod); + bind_variable_locator(py_cosserat_rod); + bind_pyelastica_interface(py_cosserat_rod); + } + } +} // namespace py_bindings + +PYBIND11_MODULE(_PyFishModel, m) { // NOLINT + m.doc() = R"pbdoc( + Bindings for Elastica++ fish model + )pbdoc"; + // Experimental : what are the drawbacks + py::module::import("elasticapp.Arrays"); + py::module::import("elasticapp.Systems.Tags"); + // + py_bindings::bind_fish_model(m); +} \ No newline at end of file diff --git a/examples/fish_swimming/FishModel/Python/CMakeLists.txt b/examples/fish_swimming/FishModel/Python/CMakeLists.txt new file mode 100644 index 00000000..89cb44bf --- /dev/null +++ b/examples/fish_swimming/FishModel/Python/CMakeLists.txt @@ -0,0 +1,20 @@ +set(LIBRARY "PyFishModel") + +elastica_python_add_module( + FishModel + MODULE_PATH "Systems/" + LIBRARY_NAME ${LIBRARY} + SOURCES + BindFishModel.cpp +) + +elastica_python_link_libraries( + ${LIBRARY} + PRIVATE + CosseratRods + pybind11::module + PyGeneratorsCosseratRod + Blaze +) + +elastica_python_add_dependencies(${LIBRARY} PyArrays PyTags) diff --git a/examples/fish_swimming/FishModel/Tags.hpp b/examples/fish_swimming/FishModel/Tags.hpp new file mode 100644 index 00000000..b098a950 --- /dev/null +++ b/examples/fish_swimming/FishModel/Tags.hpp @@ -0,0 +1,8 @@ +#pragma once + +namespace fish { + namespace tags { + struct SemiMajorAxis {}; + struct SemiMinorAxis {}; + } // namespace tags +} // namespace fish \ No newline at end of file diff --git a/examples/fish_swimming/README.md b/examples/fish_swimming/README.md new file mode 100644 index 00000000..f375455a --- /dev/null +++ b/examples/fish_swimming/README.md @@ -0,0 +1,26 @@ +## Fish swimming case study +This case study implements a custom-build elastic fish in `elasticapp` swimming in a viscous, unbounded fluid. The flow environment and fluid-structure interaction is simulated by the external solver [SophT](https://github.com/SophT-Team/SophT.git) via the immersed boundary method. Relevant publication can be found at [Tekinalp et al., 2025](https://doi.org/10.1016/j.cma.2025.117910). In particular, this study mirrors but upscales one of the demonstrations in that work. + + +### Shared memory simulation +The Python script `fish_swimming.py` showcases a single fish swimming in a smaller domain. To run this simulation, ensure you have the necessary dependencies installed: +```sh +export build_path="/path/to/elasticapp/build/dir" + +# Compile the C++ code with pybindings +cmake --build ${build_path} --target PyFishCaseStudy -j + +# Install elasticapp python bindings (e.g. into venv) +pip install /path/to/elasticapp/build/bin/python + +# Install dependencies local to this case +pip install examples/fish_swimming +``` + +### Distributed memory (MPI) simulation +This case study also supports distributed memory simulation enabled by MPI. However, a build script is not provided as the dependency management is too complex and OS-dependent. To enable MPI, the following steps should be followed + +1. Install `hdf5` with parallelization enabled. Note that the parallel build is mutually exclusive with the non-parallel version. +2. Install `h5py` that uses the parallel `hdf5`. You can verify the installation by checking the `h5py.get_config().mpi` flag. +3. Install [`sopht-mpi`](https://github.com/fankiat/sopht-mpi) along with its dependencies. This package requires Python 3.10. +4. Run the MPI fish swimming script `python fish_swimming_mpi.py --help` diff --git a/examples/fish_swimming/assets/fish_origins_32.txt b/examples/fish_swimming/assets/fish_origins_32.txt new file mode 100644 index 00000000..a97f7944 --- /dev/null +++ b/examples/fish_swimming/assets/fish_origins_32.txt @@ -0,0 +1,3 @@ +6.356402998728588827e+00,6.221866875850641243e+00,6.103384288225956134e+00,6.415019105671366795e+00,6.882145061004310449e+00,7.494702112966604268e+00,7.589019222293512001e+00,8.234124737235184455e+00,7.191766646126854567e+00,7.046384030516604113e+00,7.326142441719492027e+00,6.831360034371160239e+00,7.818514397279924921e+00,7.623007793500208074e+00,7.854811014842845829e+00,6.997312203879162418e+00,8.164777405494053397e+00,9.016315905103033046e+00,8.700516705916232496e+00,9.233355025847258091e+00,8.415711722676533668e+00,8.346114926557252289e+00,8.363950169482821551e+00,7.585607036068243936e+00,8.809133345150836902e+00,8.961481007617146943e+00,9.101009159037477758e+00,8.487360638448862460e+00,9.915949653845082423e+00,9.550891455266194541e+00,9.643718264361947234e+00,9.760295240935178640e+00 +3.641335976974819477e+00,3.205644961259203729e+00,4.624038835180411589e+00,4.454465014598867789e+00,2.404653300067767230e+00,2.281622758949836793e+00,2.808019461980479470e+00,3.650532147431748697e+00,3.557854370569322189e+00,3.354734121695098992e+00,4.178709717109614807e+00,5.192536155667914599e+00,4.366055086103965976e+00,4.783263591276819326e+00,5.948525438878473182e+00,5.619918168148214654e+00,2.061478863812614026e+00,2.395863067286029047e+00,3.739307555638875957e+00,2.742180852583353978e+00,3.217462755631639748e+00,3.268811857722347280e+00,5.117482186636629571e+00,4.717611073559889512e+00,4.456227351721122609e+00,4.657497630065422278e+00,5.577833405026642488e+00,5.741558910327813514e+00,3.430513578045028211e+00,3.596313992205511401e+00,4.339828404242206616e+00,4.842952146012654779e+00 +2.918358187349811672e+00,4.455239128553158068e+00,3.884136422089527585e+00,5.131943953643748557e+00,3.546933517961776250e+00,4.889861619191841591e+00,2.447493255385950306e+00,3.171295511474826689e+00,4.053013430590184818e+00,5.635312975195390983e+00,2.125439032750466861e+00,2.899028546353441538e+00,4.852259262139308582e+00,5.801214870688740533e+00,3.573187658462122052e+00,4.608672579107128797e+00,3.536336103705297607e+00,4.627572052597271224e+00,2.144921621598494976e+00,3.053043624013206170e+00,4.386255695845326485e+00,5.829089488293760724e+00,2.381737403189183144e+00,3.565300878547039698e+00,4.066371863096897243e+00,5.625807842238229384e+00,3.453904244409363500e+00,4.854079722146335207e+00,3.930558307163008358e+00,5.196567304195152559e+00,2.912476698004158759e+00,4.436797848044782278e+00 diff --git a/examples/fish_swimming/assets/fish_origins_96.txt b/examples/fish_swimming/assets/fish_origins_96.txt new file mode 100644 index 00000000..e7d4160c --- /dev/null +++ b/examples/fish_swimming/assets/fish_origins_96.txt @@ -0,0 +1,3 @@ +4.542783452159270041e+00,4.380946312973336809e+00,4.132390052694080218e+00,4.684490468602095170e+00,4.621350217053062259e+00,4.067682272998183635e+00,4.184109462202349583e+00,4.476009807619233172e+00,5.388640740121315886e+00,4.827422798955527838e+00,5.514497022187758724e+00,5.508357193837255750e+00,5.031817527212897545e+00,5.357640722286636681e+00,5.421778329097691085e+00,5.542713559052838690e+00,5.690175479117010937e+00,5.302520021229469194e+00,5.810151334791812339e+00,5.826764100743992536e+00,5.100928793631918090e+00,6.125973526375580569e+00,5.320134072000394454e+00,5.537641056786258176e+00,4.738107059108970454e+00,5.027218723329098360e+00,5.314586047696792548e+00,6.043977140117593017e+00,6.662526714228045321e+00,6.021536277687641103e+00,6.598546635989984033e+00,6.643669787842218177e+00,6.286863335768070549e+00,6.746738375746909711e+00,6.319954467362260075e+00,6.545378733109106228e+00,6.749343385753214797e+00,6.305058191769482789e+00,6.697561243735102643e+00,6.956087198621599299e+00,6.294776267908847700e+00,6.415457933206195840e+00,6.452311380257075513e+00,6.287353257949697571e+00,6.617899016417039704e+00,5.832639409359762439e+00,6.167890569525127376e+00,6.294054879713255168e+00,7.210811131062172308e+00,7.902092536767494124e+00,7.113202672172493912e+00,7.572608407140138276e+00,7.811571530281645082e+00,7.430870136392867309e+00,7.352877921790155646e+00,7.236427795985636457e+00,7.525374188187730518e+00,7.774493444788731722e+00,7.604890330441877566e+00,7.104441763800424248e+00,7.779160900065468454e+00,7.277883247476511386e+00,7.350405366481544078e+00,7.655632553080060632e+00,7.232952571688183063e+00,7.525401135618122517e+00,6.978413807556663961e+00,7.226953365954349451e+00,7.422919096764612590e+00,8.498719440677996673e+00,9.117847563248320242e+00,8.401669138593886643e+00,8.673194961374381506e+00,8.739939053540572544e+00,8.196890916630945156e+00,8.526355498127665911e+00,8.976581104244305109e+00,8.255768578304166994e+00,8.731941771813541919e+00,8.367112413215144429e+00,7.773112956015059360e+00,8.860032965775582525e+00,8.143942552753468789e+00,8.364376730024932982e+00,8.495524765365603415e+00,8.703214535107960614e+00,7.801671769421962388e+00,8.121297615128225900e+00,8.377342520624576849e+00,9.430911498703311580e+00,9.900099714185383704e+00,9.473213947184081718e+00,9.490497033368875179e+00,9.403199435184697919e+00,9.744505520856215242e+00,9.929911969694293461e+00,9.365677631004361459e+00,8.912801161248287585e+00,9.258853788178882382e+00 +3.765425584988125518e+00,3.082408523898994002e+00,3.568361566932360951e+00,3.391711232097598394e+00,4.798683052548090799e+00,4.140348562326009230e+00,4.643782278114122697e+00,4.471917969026447004e+00,2.349965256456878482e+00,2.675243657150152998e+00,2.588185498779356664e+00,3.398731982558944154e+00,2.863624099744036755e+00,3.737731753247147903e+00,3.242082934567967811e+00,3.897459195957469813e+00,3.240352258348639047e+00,4.581437142411339813e+00,4.235925601739113588e+00,4.820657245514403222e+00,4.304051763367285233e+00,4.850102902963840279e+00,4.279963447903880613e+00,5.499348229445732095e+00,5.284862923176516070e+00,5.426144408957771503e+00,5.150730898701481486e+00,2.459160419430980404e+00,2.070486836045351442e+00,2.140647094363047565e+00,2.448300434294446770e+00,3.097300275488257260e+00,3.308310084219185043e+00,3.976941945940669410e+00,2.946981593541631916e+00,3.052530284058791565e+00,3.181279399921091322e+00,4.215752609478986557e+00,4.860680067586131692e+00,5.019113518588809164e+00,3.950085134653193997e+00,4.022356012308513584e+00,4.253617585466322737e+00,5.283078252661224106e+00,5.877889369665765251e+00,5.842419140143597645e+00,5.762307554878765004e+00,5.215590774840070054e+00,2.389055904369501437e+00,2.150717094299980925e+00,2.011722321578134132e+00,2.219455168965805747e+00,2.997100746006094774e+00,3.130770521267711715e+00,2.985389039405428324e+00,4.031157243065021767e+00,3.390792024868264143e+00,2.821396083305305780e+00,3.952612205185202843e+00,3.992407558196809259e+00,3.989996402885652937e+00,4.977997237746082604e+00,4.404672514984016907e+00,3.847813536488946529e+00,4.937856831903517651e+00,5.579969093988682793e+00,5.989468651114965958e+00,5.690231665383632453e+00,5.004500851747367918e+00,2.657318436977283760e+00,2.639077262596605244e+00,2.245349823998288308e+00,2.667507317632052732e+00,3.726025194262068130e+00,3.865242775927439034e+00,3.254117019326729476e+00,4.051441867019037524e+00,3.288384095644641292e+00,3.512685362163004754e+00,4.825087546684111217e+00,4.759669871453708012e+00,4.316029050939661715e+00,4.868516257546531278e+00,4.269224212982698496e+00,4.537044669421561593e+00,5.363432827432536243e+00,5.886911684892975671e+00,5.806018452831225751e+00,5.371711746605489779e+00,3.377726265047564347e+00,3.596350538088410875e+00,2.911855978482955187e+00,3.462284237710037793e+00,4.476633189462133622e+00,4.760641553862984487e+00,4.144173262300984817e+00,4.521598044321639698e+00,5.527493327279904278e+00,5.214359980940487027e+00 +2.314419387595163524e+00,3.500308890405597673e+00,4.600638302356595766e+00,5.676509929875433791e+00,2.618322117248852887e+00,3.400865097985482421e+00,4.381751739203862961e+00,5.460036847894836676e+00,3.468276371583518536e+00,4.584291230997976641e+00,5.520737148456400512e+00,1.557486268553734954e+00,2.508275899774045303e+00,3.287633941703477447e+00,4.367675786558592677e+00,5.299672408632204679e+00,6.447229012746931609e+00,1.683871032447608274e+00,2.404908969244566208e+00,3.518769029730313580e+00,4.274274969555730053e+00,4.814445892874811506e+00,6.450451649739030735e+00,2.654963883937600855e+00,3.582166284253114341e+00,4.732684601256615053e+00,5.784035470156774217e+00,2.340861481518110576e+00,3.282208115930608461e+00,4.520471409508677318e+00,5.851067598403679071e+00,1.346008550227505651e+00,2.515956459438009052e+00,3.314901667598512791e+00,3.638219835563577309e+00,4.973712671415550624e+00,6.726104433023717633e+00,1.099269404033211384e+00,2.637814550428239091e+00,3.858825796257719265e+00,4.261147285127486306e+00,5.739188323413210391e+00,6.925484161359449331e+00,1.810272149878307202e+00,3.038082368447215575e+00,3.936679100988996538e+00,5.152310760508933996e+00,6.275709609357598140e+00,2.232618396825466700e+00,3.294453545258468452e+00,4.314685298547860093e+00,5.241761735364840469e+00,1.533943004403326427e+00,2.893371880788886408e+00,4.093695120209703830e+00,4.698858146489801157e+00,5.643195593702616364e+00,6.298711940668543008e+00,1.061835509409807976e+00,2.069203138975949319e+00,3.738734355974626222e+00,5.013826575046608447e+00,5.923024464170089765e+00,6.919182404477562720e+00,1.359737110177305031e+00,2.234933967102876284e+00,4.317431416851846393e+00,5.590668020971590657e+00,6.561279180080472173e+00,2.356148255988975482e+00,3.401818630556790346e+00,4.338434093656736046e+00,5.487179377188774687e+00,1.589876717579238097e+00,2.446345708805603092e+00,3.525827553938717607e+00,4.366239250691846152e+00,4.745908020032960195e+00,6.338357641947657939e+00,1.633242296079842948e+00,2.918458460728281967e+00,3.264416051806617869e+00,4.185456044335710502e+00,5.283798296813361439e+00,6.473684087644921981e+00,2.613562550994360301e+00,3.407892651090386860e+00,4.641628747895832241e+00,5.695028905582493195e+00,2.509249106565001775e+00,3.525444163315542223e+00,4.471854318234813164e+00,5.465939886979597873e+00,2.351384104839897127e+00,3.585528195804669416e+00,4.611442340607201373e+00,5.672035464589169074e+00,3.688406074149350289e+00,4.763270649325708206e+00 diff --git a/examples/fish_swimming/fish_swimming.py b/examples/fish_swimming/fish_swimming.py new file mode 100644 index 00000000..79d34235 --- /dev/null +++ b/examples/fish_swimming/fish_swimming.py @@ -0,0 +1,178 @@ +import numpy as np + +import sopht.simulator as sps + +from elasticapp.Systems import FishModel +from elasticapp.CaseStudies.PyFishCaseStudy import FishCaseStudy +from elasticapp.Arrays._PyArrays import Matrix + +import fish_utils as fu + + +def run_fish_swimming( + non_dim_final_time: float, + n_elem: int, + grid_size_x: int, + surface_grid_density_for_largest_element: int, + slenderness_ratio: float, + mass_ratio: float, + non_dim_bending_stiffness: float, + actuation_reynolds_number: float, + coupling_stiffness: float = -3.2e4, + coupling_damping: float = -32.0, + flow_num_threads: int = 4, + save_data: bool = False, + data_dir: str | None = None, +) -> None: + # ================ Global parameters ================ + grid_dim = 3 + rho_f = 1.0 + base_length = 1.0 + x_range, y_range, z_range = (3.5, 1.0, 0.8) + dx = x_range / grid_size_x + grid_size_y = int(np.ceil(y_range / dx)) + grid_size_z = int(np.ceil(z_range / dx)) + + # ================ Fish parameters ================ + head_position = Matrix(np.array([[2.3, 0.5, 0.4]], dtype=np.float64).T) + period = 1.0 + vel_scale = base_length / period + rho_s = mass_ratio * rho_f + base_diameter = base_length / slenderness_ratio + base_radius = base_diameter / 2 + moment_of_inertia = np.pi / 4 * base_radius**4 + youngs_modulus = ( + non_dim_bending_stiffness + * rho_f + * vel_scale**2 + * base_length**3 + * base_diameter + / moment_of_inertia + ) + + # ================ Fish simulator ================ + cs = FishCaseStudy( + num_fish=1, + origins=head_position, + n_elements=n_elem, + density=rho_s, + length=base_length, + damping_rate=0.16, + youngs_modulus=youngs_modulus, + dt_scale=5e-4, + period=period, + ) + elastic_fish_list, virtual_fish_list, flow_forces, flow_torques = cs.communicate() + flow_forces = np.asarray(flow_forces) + flow_torques = np.asarray(flow_torques) + + # ================ Flow simulator ================ + kinematic_viscosity = base_length * vel_scale / actuation_reynolds_number + flow_sim = sps.create_unbounded_flow_simulator_3d( + grid_size=(grid_size_z, grid_size_y, grid_size_x), + x_range=x_range, + kinematic_viscosity=kinematic_viscosity, + flow_type="navier_stokes_with_forcing", + real_t=np.float64, + num_threads=flow_num_threads, + filter_vorticity=True, + filter_setting_dict={"order": 5, "type": "convolution"}, + ) + + # ============== Immersed boundary =============== + flow_interaction = fu.FishFlowInteraction( + cosserat_rod=elastic_fish_list[0], + eul_grid_forcing_field=flow_sim.eul_grid_forcing_field, + eul_grid_velocity_field=flow_sim.velocity_field, + virtual_boundary_stiffness_coeff=coupling_stiffness, + virtual_boundary_damping_coeff=coupling_damping, + dx=flow_sim.dx, + grid_dim=grid_dim, + real_t=np.float64, + num_threads=flow_num_threads, + forcing_grid_cls=fu.FishSurfaceForcingGrid, + surface_grid_density_for_largest_element=surface_grid_density_for_largest_element, + ) + + # ================ IO ================ + if save_data: + data_dir = fu.get_data_path(data_dir) + flow_data_manager = fu.ImageDataIO(dx=flow_sim.dx, grid_size=flow_sim.grid_size) + fish_data_manager = fu.FishIO(elastic_fish_list, virtual_fish_list) + grid_data_manager = fu.GridIO( + num_fish=1, num_lag_points_per_fish=flow_interaction.forcing_grid.num_lag_nodes + ) + + # ================ Simulation ================ + final_time = non_dim_final_time * period + data_timer_limit = period / 20.0 + data_timer = 0.0 + rod_dt = cs.current_dt() + + while flow_sim.time < final_time: + print(f"{flow_sim.time / final_time * 100:.2f}%", end="\r") + # Save data + if save_data and (data_timer == 0.0 or data_timer > data_timer_limit): + data_timer = 0.0 + flow_data_manager.write_to_file( + filename=data_dir / f"flow_{int(flow_sim.time * 1000):04d}.vtkhdf", + time=flow_sim.time, + velocity=flow_sim.velocity_field, + vorticity=flow_sim.vorticity_field, + ) + fish_data_manager.write_to_file( + filename=data_dir / f"fish_{int(cs.current_time() * 1000):04d}.h5", + time=np.float64(cs.current_time()), + ) + grid_data_manager.write_to_file( + filename=data_dir / f"grid_{int(cs.current_time() * 1000):04d}.vtkhdf", + time=np.float64(cs.current_time()), + position=flow_interaction.forcing_grid.position_field, + velocity=flow_interaction.forcing_grid.velocity_field, + ) + + # Compute time steps + flow_dt = min(flow_sim.compute_stable_timestep(dt_prefac=0.25), rod_dt * 20) + rod_time_steps = max(int(flow_dt / rod_dt), 1) + local_rod_dt = flow_dt / rod_time_steps + + # Rod + IBM steps + for _ in range(rod_time_steps): + flow_interaction.compute_flow_forces_and_torques() + flow_interaction.time_step(dt=local_rod_dt) + + # Copy flow forces and torques into C++ buffer + flow_forces[...] = flow_interaction.body_flow_forces + flow_torques[...] = flow_interaction.body_flow_torques + + cs.step_with(dt=local_rod_dt) + + # Update IBM grid + flow_interaction() + + # Flow step + flow_sim.time_step(dt=flow_dt) + + # Update data timer + data_timer += flow_dt + + +if __name__ == "__main__": + import logging + logging.basicConfig(level=logging.INFO) + + nx = 96 + n_elements = nx // 4 + n_max = n_elements // 4 + run_fish_swimming( + non_dim_final_time=0.06, + n_elem=n_elements, + grid_size_x=nx, + surface_grid_density_for_largest_element=n_max, + slenderness_ratio=12.500535121608523, + mass_ratio=1.0, + non_dim_bending_stiffness=0.565414058870102, + actuation_reynolds_number=7142.857142857143, + save_data=True, + data_dir="fish_swimming_data", + ) diff --git a/examples/fish_swimming/fish_swimming_mpi.py b/examples/fish_swimming/fish_swimming_mpi.py new file mode 100644 index 00000000..2196fdb7 --- /dev/null +++ b/examples/fish_swimming/fish_swimming_mpi.py @@ -0,0 +1,479 @@ +import os +import click +from time import time +from typing import Optional + +import numpy as np +from mpi4py import MPI +from sopht.utils import get_real_t +from sopht_mpi.simulator import UnboundedFlowSimulator3D +from sopht_mpi.simulator.immersed_body import ImmersedBodyFlowInteractionMPI + +import fish_utils as fu + +from elasticapp.Systems import FishModel # noqa: F401 +from elasticapp.CaseStudies.PyFishCaseStudy import FishCaseStudy +from elasticapp.Arrays._PyArrays import Matrix + + +def reload_flow_interaction(flow_interaction: ImmersedBodyFlowInteractionMPI, time: float) -> None: + """Reload the flow interaction state from the given time. + + Must be called after restart. + """ + + flow_interaction.time = time + flow_interaction.forcing_grid.compute_lag_grid_position_field() + flow_interaction.forcing_grid.compute_lag_grid_velocity_field() + flow_interaction.mpi_lagrangian_field_communicator.map_lagrangian_nodes_based_on_position( + flow_interaction.forcing_grid.position_field + ) + flow_interaction.local_num_lag_nodes = ( + flow_interaction.mpi_lagrangian_field_communicator.local_num_lag_nodes + ) + flow_interaction._init_local_buffers(flow_interaction.local_num_lag_nodes) + flow_interaction.mpi_lagrangian_field_communicator.scatter_global_field( + local_lag_field=flow_interaction.local_lag_grid_position_mismatch_field, + global_lag_field=flow_interaction.global_lag_grid_position_mismatch_field, + ) + flow_interaction.mpi_lagrangian_field_communicator.scatter_global_field( + local_lag_field=flow_interaction.local_lag_grid_velocity_mismatch_field, + global_lag_field=flow_interaction.global_lag_grid_velocity_mismatch_field, + ) + + +def distribute_fish( + mpi_size: int, + mpi_rank: int, + num_fish_total: int, + mode: str = "kz", +) -> tuple[list, list]: + """Distribute the fish across MPI ranks based on the given mode: + + even: Distribute the fish as evenly as possible across all ranks. + kz: Allocate all fish to rank 0. + """ + if mode.lower() == "even": + num_fish_per_rank = np.ones((mpi_size,), np.int32) * (num_fish_total // mpi_size) + rem = num_fish_total % mpi_size + num_fish_per_rank[:rem] += 1 + indices = np.insert(np.cumsum(num_fish_per_rank), 0, 0) + return num_fish_per_rank, list(range(indices[mpi_rank], indices[mpi_rank + 1])) + elif mode.lower() == "kz": + num_fish_per_rank = [0] * mpi_size + num_fish_per_rank[0] = num_fish_total + index = list(range(num_fish_total)) if mpi_rank == 0 else [] + return num_fish_per_rank, index + else: + raise ValueError("Invalid mode. Use 'even' or 'kz'.") + + +@click.command() +@click.option( + "--non-dim-final-time", + "-t", + type=float, + default=1.0, + help="Non-dimensional final time of the simulation. Scaled by the period.", +) +@click.option("--n-elem", "-n", type=int, default=16, help="Number of elements per fish.") +@click.option( + "--grid-size-x", + "-g", + type=int, + default=192, + help="Number of flow grid points in the x-direction.", +) +@click.option( + "--surface-grid-density-for-largest-element", + "-s", + type=int, + default=4, + help="Surface grid density for the largest element.", +) +@click.option( + "--slenderness-ratio", + "-r", + type=float, + default=12.5, + help="Slenderness ratio of the fish.", +) +@click.option( + "--mass_ratio", + "-m", + type=float, + default=1.0, + help="Mass ratio of the fish to the fluid.", +) +@click.option( + "--non-dim-bending-stiffness", + "-b", + type=float, + default=0.565414058870102, + help="Non-dimensional bending stiffness of the fish.", +) +@click.option( + "--actuation-reynolds-number", + "-a", + type=float, + default=7142.857142857143, + help="Actuation Reynolds number for the fish.", +) +@click.option( + "--coupling-stiffness", + "-k", + type=float, + default=-32000.0, + help="Coupling stiffness for the immersed boundary method.", +) +@click.option( + "--coupling-damping", + "-c", + type=float, + default=-32.0, + help="Coupling damping for the immersed boundary method.", +) +@click.option( + "--precision", + "-p", + type=str, + default="double", + help="Precision of the simulation (single or double).", +) +@click.option( + "--rank-distribution", + "-R", + nargs=3, + type=(int, int, int), + default=None, + help="Rank distribution for MPI.", +) +@click.option( + "--distribution-mode", + "-D", + type=click.Choice(["even", "kz"], case_sensitive=False), + default="kz", + help="Distribution mode for the fish across MPI ranks.", +) +@click.option( + "--restart/--no-restart", + default=False, + help="Flag to restart the simulation from the last saved data in data_dir.", +) +@click.option( + "--data-dir", + "-d", + type=str, + default=None, + help="Directory to save (and load) the simulation data.", +) +@click.option( + "--fish-head-positions-file", + "-h", + type=str, + default="assets/fish_head_positions.txt", + help="File containing the fish head positions.", +) +def run_fish_swimming_mpi( + non_dim_final_time: float, + n_elem: int, + grid_size_x: int, + surface_grid_density_for_largest_element: int, + slenderness_ratio: float, + mass_ratio: float, + non_dim_bending_stiffness: float, + actuation_reynolds_number: float, + coupling_stiffness: float = -3.2e4, + coupling_damping: float = -32.0, + precision: str = "double", + rank_distribution: Optional[tuple[int, int, int]] = None, + distribution_mode: str = "kz", + restart: bool = False, + data_dir: str | None = None, + fish_head_positions_file: str = "assets/fish_head_positions.txt", +) -> None: + # ================ Global parameters ================ + grid_dim = 3 + real_t = get_real_t(precision) + rho_f = 1.0 + base_length = 1.0 + x_range, y_range, z_range = (12.0, 8.0, 8.0) + dx = x_range / grid_size_x + grid_size_y = int(np.ceil(y_range / dx)) + grid_size_z = int(np.ceil(z_range / dx)) + + # ================ Fish parameters ================ + period = 1.0 + vel_scale = base_length / period + rho_s = mass_ratio * rho_f + base_diameter = base_length / slenderness_ratio + base_radius = base_diameter / 2 + moment_of_inertia = np.pi / 4 * base_radius**4 + youngs_modulus = ( + non_dim_bending_stiffness + * rho_f + * vel_scale**2 + * base_length**3 + * base_diameter + / moment_of_inertia + ) + origins = np.loadtxt(fish_head_positions_file, delimiter=",") + + # ================ Flow simulator ================ + kinematic_viscosity = base_length * vel_scale / actuation_reynolds_number + flow_sim = UnboundedFlowSimulator3D( + grid_size=(grid_size_z, grid_size_y, grid_size_x), + x_range=x_range, + kinematic_viscosity=kinematic_viscosity, + flow_type="navier_stokes_with_forcing", + real_t=real_t, + rank_distribution=rank_distribution, + filter_vorticity=True, + filter_setting_dict={"order": 5, "type": "convolution"}, + ) + mpi_rank = flow_sim.mpi_construct.rank + mpi_size = flow_sim.mpi_construct.size + + # Get the data directory and broadcast + if mpi_rank == 0: + data_dir = fu.get_data_path(data_dir) + else: + data_dir = None + + data_dir = flow_sim.mpi_construct.grid.bcast(data_dir, root=0) + + # ================ Fish simulator ================ + # Distribute fishes evenly across MPI ranks + num_fish_total = origins.shape[1] + num_fish_per_rank, index = distribute_fish( + mpi_size=mpi_size, + mpi_rank=mpi_rank, + num_fish_total=num_fish_total, + mode=distribution_mode, + ) + if mpi_size == 1: + origin = Matrix(origins) + else: + tmp = origins[:, index].copy() + origin = Matrix(tmp) + + has_fish = num_fish_per_rank[mpi_rank] > 0 + num_proc_with_fish = np.count_nonzero(num_fish_per_rank) + + # Initialize the case study only if there are fishes on this rank + if has_fish: + cs = FishCaseStudy( + num_fish=int(num_fish_per_rank[mpi_rank]), + origins=origin, + n_elements=n_elem, + density=rho_s, + length=base_length, + damping_rate=0.16, + youngs_modulus=youngs_modulus, + dt_scale=5e-4, + period=period, + ) + elastic_fish_list, virtual_fish_list, flow_forces, flow_torques = cs.communicate() + flow_forces = np.asarray(flow_forces) + flow_torques = np.asarray(flow_torques) + else: + elastic_fish_list = [] + virtual_fish_list = [] + flow_forces = None + flow_torques = None + + # ============== Immersed boundary =============== + flow_interactions = [] + + # Each rank has flow_interaction objects for all fishes + # But the only functional ones are those that corresponds to the fish + # on this rank. + auto_ghosting = True + for rank_idx in range(num_proc_with_fish): + flow_interactions.append( + fu.MultiFishFlowInteractionMPI( + mpi_construct=flow_sim.mpi_construct, + mpi_ghost_exchange_communicator=flow_sim.mpi_ghost_exchange_communicator, + fish_list=elastic_fish_list, + num_fish=num_fish_per_rank[rank_idx], + n_elem_per_fish=n_elem, + eul_grid_forcing_field=flow_sim.eul_grid_forcing_field, + eul_grid_velocity_field=flow_sim.velocity_field, + virtual_boundary_stiffness_coeff=coupling_stiffness, + virtual_boundary_damping_coeff=coupling_damping, + dx=flow_sim.dx, + grid_dim=grid_dim, + forcing_grid_cls=fu.MultiFishForcingGrid, + surface_grid_density_for_largest_element=surface_grid_density_for_largest_element, + master_rank=rank_idx, + auto_ghosting=auto_ghosting, + ) + ) + auto_ghosting = False + for idx, flow_interaction in enumerate(flow_interactions): + if idx == mpi_rank: + assert not flow_interaction.is_empty() + else: + assert flow_interaction.is_empty() + + print(f"Rank: {mpi_rank}, Number of fish: {num_fish_per_rank[mpi_rank]}") + + # ================ IO ================ + if has_fish: + num_lag_node_per_fish = ( + flow_interactions[mpi_rank].forcing_grid.num_lag_nodes // num_fish_per_rank[mpi_rank] + ) + lag_fields = { + "position_mismatch": flow_interactions[ + mpi_rank + ].global_lag_grid_position_mismatch_field, + "velocity_mismatch": flow_interactions[ + mpi_rank + ].global_lag_grid_velocity_mismatch_field, + } + else: + num_lag_node_per_fish = 0 + lag_fields = { + "position_mismatch": None, + "velocity_mismatch": None, + } + + num_lag_node_per_fish = flow_sim.mpi_construct.grid.bcast(num_lag_node_per_fish, root=0) + fish_writer = fu.MPIFishIO( + mpi_comm=flow_sim.mpi_construct.grid, + fish_list=elastic_fish_list, + virtual_fish_list=virtual_fish_list, + index_list=index, + num_fish_element=n_elem, + num_lag_nodes=num_lag_node_per_fish, + global_num_fish=num_fish_total, + **lag_fields, + ) + + flow_writer = fu.MPIImageDataIO( + mpi_comm=flow_sim.mpi_construct.grid, + dx=flow_sim.dx, + global_grid_size=flow_sim.grid_size, + local_grid_size=flow_sim.mpi_construct.local_grid_size, + ghost_size=flow_sim.ghost_size, + ) + + if restart: + fish_file = fu.get_most_recent_data_file(data_dir, "fish_*.h5") + flow_file = fu.get_most_recent_data_file(data_dir, "flow_*.vtkhdf") + fish_time = fish_writer.read_from_file(fish_file) + flow_sim.time = flow_writer.read_from_file( + flow_file, + velocity=flow_sim.velocity_field, + vorticity=flow_sim.vorticity_field, + ) + if has_fish: + cs.set_current_time(fish_time) + if mpi_rank == 0: + print("Loaded fish data from ", fish_file) + print("Loaded flow data from ", flow_file) + + for flow_interaction in flow_interactions: + reload_flow_interaction(flow_interaction, fish_time) + + # ================ Simulation ================ + final_time = non_dim_final_time * period + data_timer_limit = period / 20.0 + data_timer = 0.0 + if has_fish: + rod_dt = cs.current_dt() + else: + rod_dt = None + + rod_dt = flow_sim.mpi_construct.grid.bcast(rod_dt, root=0) + + flow_time = 0.0 + rod_time = 0.0 + fsi_time = 0.0 + io_time = 0.0 + + start_time = time() + while flow_sim.time < final_time: + if data_timer == 0.0 or data_timer > data_timer_limit: + data_timer = 0.0 + + # Print progress only on rank 0 + if mpi_rank == 0: + # print(f"{flow_sim.time / final_time * 100:.2f}%", flush=True) + print(flow_sim.time, io_time, flow_time, rod_time, fsi_time, flush=True) + + a = time() + + # Log fish data (Not VTKHDF. Post-processing required) + fish_time = cs.current_time() if has_fish else 0.0 + fish_time = flow_sim.mpi_construct.grid.bcast(fish_time, root=0) + fish_writer.write_to_file( + filename=os.path.join(data_dir, f"fish_{int(flow_sim.time * 1000):04d}.h5"), + time=fish_time, + ) + + # Log flow data for all ranks (HDF5-MPI) + flow_writer.write_to_file( + filename=os.path.join(data_dir, f"flow_{int(flow_sim.time * 1000):04d}.vtkhdf"), + time=flow_sim.time, + velocity=flow_sim.velocity_field, + vorticity=flow_sim.vorticity_field, + ) + io_time += time() - a + + # Compute time steps + flow_dt = min( + flow_sim.compute_stable_timestep(dt_prefac=0.25, precision=precision), + rod_dt * 100, + ) + rod_time_steps = max(int(flow_dt / rod_dt), 1) + local_rod_dt = flow_dt / rod_time_steps + + # Rod + IBM steps + for _ in range(rod_time_steps): + a = time() + for flow_interaction in flow_interactions: + flow_interaction.compute_flow_forces_and_torques() + if has_fish: + flow_forces[...] = flow_interactions[mpi_rank].body_flow_forces + flow_torques[...] = flow_interactions[mpi_rank].body_flow_torques + flow_interaction.time_step(dt=local_rod_dt) + fsi_time += time() - a + + if has_fish: + a = time() + cs.step_with(updated_dt=local_rod_dt) + rod_time += time() - a + + # Update IBM grid + a = time() + for flow_interaction in flow_interactions: + flow_interaction() + flow_interaction.update_buffers(flow_interaction.forcing_grid.position_field) + fsi_time += time() - a + + # Flow step + a = time() + flow_sim.time_step(dt=flow_dt, free_stream_velocity=None) + flow_time += time() - a + + # Update data timer + data_timer += flow_dt + + flow_time = flow_sim.mpi_construct.grid.allreduce(flow_time, op=MPI.MAX) + rod_time = flow_sim.mpi_construct.grid.allreduce(rod_time, op=MPI.MAX) + fsi_time = flow_sim.mpi_construct.grid.allreduce(fsi_time, op=MPI.MAX) + io_time = flow_sim.mpi_construct.grid.allreduce(io_time, op=MPI.MAX) + + # Finalize the simulation + sim_time = time() - start_time + + sim_time = flow_sim.mpi_construct.grid.reduce(sim_time, op=MPI.MAX, root=0) + if mpi_rank == 0: + print(f"Simulation completed in {sim_time:.2f} seconds.") + + MPI.Finalize() + + +if __name__ == "__main__": + run_fish_swimming_mpi() diff --git a/examples/fish_swimming/pyproject.toml b/examples/fish_swimming/pyproject.toml new file mode 100644 index 00000000..2f0b7a43 --- /dev/null +++ b/examples/fish_swimming/pyproject.toml @@ -0,0 +1,16 @@ +[build-system] +requires = ["setuptools", "pybind11", "numpy"] +build-backend = "setuptools.build_meta" + +[project] +name = "fish_swimming_utils" +version = "0.0.1" +description = "Python utilities for fish swimming simulations" +requires-python = ">=3.10, <3.14" +dependencies = [ + "h5py", + "sopht @ git+https://github.com/SophT-Team/SophT.git@v0.0.3", +] + +[tool.setuptools.package-dir] +"fish_utils" = "utils" diff --git a/examples/fish_swimming/setup.py b/examples/fish_swimming/setup.py new file mode 100644 index 00000000..e11b3d95 --- /dev/null +++ b/examples/fish_swimming/setup.py @@ -0,0 +1,13 @@ +import numpy +import pybind11 +from setuptools import Extension, setup + +extension = Extension( + name="fish_utils._fish_grid_helper", + sources=["utils/fish_grid.cpp"], + include_dirs=[pybind11.get_include(), numpy.get_include()], + language="c++", + extra_compile_args=["-O3", "-std=c++20", "-march=native"], +) + +setup(ext_modules=[extension]) diff --git a/examples/fish_swimming/utils/__init__.py b/examples/fish_swimming/utils/__init__.py new file mode 100644 index 00000000..b0de87e9 --- /dev/null +++ b/examples/fish_swimming/utils/__init__.py @@ -0,0 +1,39 @@ +from ._fish_grid_helper import FishGridHelper +from .fish_flow_interaction import FishFlowInteraction, MultiFishFlowInteraction +from .fish_grid import ( + FishSurfaceForcingGrid, + MultiFishForcingGrid, + MultiFishSurfaceForcingGrid, +) +from .fish_io import FishIO, GridIO, ImageDataIO +from .path_utils import get_data_path, get_most_recent_data_file + +try: + from .fish_flow_interaction_mpi import ( + FishFlowInteractionMPI, + MultiFishFlowInteractionMPI, + ) + from .fish_io_mpi import MPIFishIO, MPIImageDataIO + + __all__ = [ + "FishFlowInteractionMPI", + "MultiFishFlowInteractionMPI", + "MPIImageDataIO", + "MPIFishIO", + ] +except ImportError: + __all__ = [] + +__all__ += [ + "FishSurfaceForcingGrid", + "MultiFishSurfaceForcingGrid", + "MultiFishForcingGrid", + "FishFlowInteraction", + "MultiFishFlowInteraction", + "FishGridHelper", + "get_data_path", + "get_most_recent_data_file", + "ImageDataIO", + "FishIO", + "GridIO", +] diff --git a/examples/fish_swimming/utils/fish_flow_interaction.py b/examples/fish_swimming/utils/fish_flow_interaction.py new file mode 100644 index 00000000..c9ca1ef4 --- /dev/null +++ b/examples/fish_swimming/utils/fish_flow_interaction.py @@ -0,0 +1,109 @@ +import numpy as np +from sopht.simulator.immersed_body import ( + ImmersedBodyForcingGrid, + ImmersedBodyFlowInteraction, +) +from elasticapp.Systems.FishModel._PyFishModel import Fish as RodType +from typing import Type, Optional + + +class FishFlowInteraction(ImmersedBodyFlowInteraction): + """Elastica++ flow interaction construct for a single fish.""" + + def __init__( + self, + cosserat_rod: RodType, + eul_grid_forcing_field: np.ndarray, + eul_grid_velocity_field: np.ndarray, + virtual_boundary_stiffness_coeff: float, + virtual_boundary_damping_coeff: float, + dx: float, + grid_dim: int, + forcing_grid_cls: Type[ImmersedBodyForcingGrid], + real_t: type = np.float64, + eul_grid_coord_shift: Optional[float] = None, + interp_kernel_width: Optional[float] = None, + enable_eul_grid_forcing_reset: bool = False, + num_threads: int | bool = False, + start_time: float = 0.0, + **forcing_grid_kwargs, + ) -> None: + n_elems = int(cosserat_rod.n_elems[0]) + body_flow_forces = np.zeros( + (3, n_elems + 1), + ) + body_flow_torques = np.zeros( + (3, n_elems), + ) + + forcing_grid_kwargs["cosserat_rod"] = cosserat_rod + + super().__init__( + eul_grid_forcing_field, + eul_grid_velocity_field, + body_flow_forces, + body_flow_torques, + forcing_grid_cls, + virtual_boundary_stiffness_coeff, + virtual_boundary_damping_coeff, + dx, + grid_dim, + real_t, + eul_grid_coord_shift, + interp_kernel_width, + enable_eul_grid_forcing_reset, + num_threads, + start_time, + **forcing_grid_kwargs, + ) + + +class MultiFishFlowInteraction(ImmersedBodyFlowInteraction): + """Elastica++ flow interaction construct for multiple fish.""" + + def __init__( + self, + fish_list: list[RodType], + eul_grid_forcing_field: np.ndarray, + eul_grid_velocity_field: np.ndarray, + virtual_boundary_stiffness_coeff: float, + virtual_boundary_damping_coeff: float, + dx: float, + grid_dim: int, + forcing_grid_cls: Type[ImmersedBodyForcingGrid], + real_t: type = np.float64, + eul_grid_coord_shift: Optional[float] = None, + interp_kernel_width: Optional[float] = None, + enable_eul_grid_forcing_reset: bool = False, + num_threads: int | bool = False, + start_time: float = 0.0, + **forcing_grid_kwargs, + ) -> None: + n_elems = n_nodes = 0 + for fish in fish_list: + n_elems += int(fish.n_elems[0]) + n_nodes += int(fish.n_elems[0]) + 1 + + body_flow_forces = np.zeros((3, n_nodes), dtype=real_t) + body_flow_torques = np.zeros((3, n_elems), dtype=real_t) + + forcing_grid_kwargs["fish_list"] = fish_list + + super().__init__( + eul_grid_forcing_field, + eul_grid_velocity_field, + body_flow_forces, + body_flow_torques, + forcing_grid_cls, + virtual_boundary_stiffness_coeff, + virtual_boundary_damping_coeff, + dx, + grid_dim, + real_t, + eul_grid_coord_shift, + interp_kernel_width, + enable_eul_grid_forcing_reset, + num_threads, + start_time, + **forcing_grid_kwargs, + ) diff --git a/examples/fish_swimming/utils/fish_flow_interaction_mpi.py b/examples/fish_swimming/utils/fish_flow_interaction_mpi.py new file mode 100644 index 00000000..8ed3c67f --- /dev/null +++ b/examples/fish_swimming/utils/fish_flow_interaction_mpi.py @@ -0,0 +1,147 @@ +import numpy as np +from sopht.simulator.immersed_body import ( + ImmersedBodyForcingGrid, +) +from sopht_mpi.simulator.immersed_body import ( + EmptyForcingGrid, + ImmersedBodyFlowInteractionMPI, +) +from sopht_mpi.utils import ( + MPIConstruct2D, + MPIConstruct3D, + MPIGhostCommunicator2D, + MPIGhostCommunicator3D, +) +from elasticapp.Systems.FishModel._PyFishModel import Fish as RodType +from typing import Type, Optional + + +class FishFlowInteractionMPI(ImmersedBodyFlowInteractionMPI): + """Class for Elastica++ Cosserat rod flow interaction.""" + + def __init__( + self, + mpi_construct: "MPIConstruct2D | MPIConstruct3D", + mpi_ghost_exchange_communicator: "MPIGhostCommunicator2D | MPIGhostCommunicator3D", + cosserat_rod: RodType | int, + eul_grid_forcing_field: np.ndarray, + eul_grid_velocity_field: np.ndarray, + virtual_boundary_stiffness_coeff: float, + virtual_boundary_damping_coeff: float, + dx: float, + grid_dim: int, + forcing_grid_cls: Type[ImmersedBodyForcingGrid], + eul_grid_coord_shift: Optional[float] = None, + interp_kernel_width: Optional[float] = None, + enable_eul_grid_forcing_reset: bool = False, + start_time: float = 0.0, + master_rank: int = 0, + assume_data_locality: bool = False, + auto_ghosting: bool = True, + **forcing_grid_kwargs, + ) -> None: + self.master_rank = master_rank + if mpi_construct.rank == self.master_rank: + if forcing_grid_cls is None: + raise ValueError( + "forcing_grid_cls must be provided when master_rank is 0" + ) + self.forcing_grid = forcing_grid_cls( + grid_dim=grid_dim, + cosserat_rod=cosserat_rod, + **forcing_grid_kwargs, + ) + n_elems = int(cosserat_rod.n_elems[0]) + self.body_flow_forces = np.zeros((3, n_elems + 1)) + self.body_flow_torques = np.zeros((3, n_elems)) + else: + self.forcing_grid = EmptyForcingGrid(grid_dim=grid_dim) + self.body_flow_forces = np.zeros((3, cosserat_rod + 1)) + self.body_flow_torques = np.zeros((3, cosserat_rod)) + + # initialising super class + super().__init__( + mpi_construct=mpi_construct, + mpi_ghost_exchange_communicator=mpi_ghost_exchange_communicator, + eul_grid_forcing_field=eul_grid_forcing_field, + eul_grid_velocity_field=eul_grid_velocity_field, + virtual_boundary_stiffness_coeff=virtual_boundary_stiffness_coeff, + virtual_boundary_damping_coeff=virtual_boundary_damping_coeff, + dx=dx, + grid_dim=grid_dim, + eul_grid_coord_shift=eul_grid_coord_shift, + interp_kernel_width=interp_kernel_width, + enable_eul_grid_forcing_reset=enable_eul_grid_forcing_reset, + start_time=start_time, + assume_data_locality=assume_data_locality, + auto_ghosting=auto_ghosting, + ) + + def is_empty(self) -> bool: + return isinstance(self.forcing_grid, EmptyForcingGrid) + + +class MultiFishFlowInteractionMPI(ImmersedBodyFlowInteractionMPI): + """Class for Elastica++ Cosserat rod flow interaction.""" + + def __init__( + self, + mpi_construct: "MPIConstruct2D | MPIConstruct3D", + mpi_ghost_exchange_communicator: "MPIGhostCommunicator2D | MPIGhostCommunicator3D", + fish_list: list[RodType], + num_fish: int, + n_elem_per_fish: int, + eul_grid_forcing_field: np.ndarray, + eul_grid_velocity_field: np.ndarray, + virtual_boundary_stiffness_coeff: float, + virtual_boundary_damping_coeff: float, + dx: float, + grid_dim: int, + forcing_grid_cls: Type[ImmersedBodyForcingGrid], + eul_grid_coord_shift: Optional[float] = None, + interp_kernel_width: Optional[float] = None, + enable_eul_grid_forcing_reset: bool = False, + start_time: float = 0.0, + master_rank: int = 0, + assume_data_locality: bool = False, + auto_ghosting: bool = True, + **forcing_grid_kwargs, + ) -> None: + self.master_rank = master_rank + if mpi_construct.rank == self.master_rank: + if forcing_grid_cls is None: + raise ValueError( + "forcing_grid_cls must be provided when master_rank is 0" + ) + self.forcing_grid = forcing_grid_cls( + grid_dim=grid_dim, + fish_list=fish_list, + **forcing_grid_kwargs, + ) + + else: + self.forcing_grid = EmptyForcingGrid(grid_dim=grid_dim) + + self.body_flow_forces = np.zeros((3, (n_elem_per_fish + 1) * num_fish)) + self.body_flow_torques = np.zeros((3, n_elem_per_fish * num_fish)) + + # initialising super class + super().__init__( + mpi_construct=mpi_construct, + mpi_ghost_exchange_communicator=mpi_ghost_exchange_communicator, + eul_grid_forcing_field=eul_grid_forcing_field, + eul_grid_velocity_field=eul_grid_velocity_field, + virtual_boundary_stiffness_coeff=virtual_boundary_stiffness_coeff, + virtual_boundary_damping_coeff=virtual_boundary_damping_coeff, + dx=dx, + grid_dim=grid_dim, + eul_grid_coord_shift=eul_grid_coord_shift, + interp_kernel_width=interp_kernel_width, + enable_eul_grid_forcing_reset=enable_eul_grid_forcing_reset, + start_time=start_time, + assume_data_locality=assume_data_locality, + auto_ghosting=auto_ghosting, + ) + + def is_empty(self) -> bool: + return isinstance(self.forcing_grid, EmptyForcingGrid) diff --git a/examples/fish_swimming/utils/fish_grid.cpp b/examples/fish_swimming/utils/fish_grid.cpp new file mode 100644 index 00000000..7980e264 --- /dev/null +++ b/examples/fish_swimming/utils/fish_grid.cpp @@ -0,0 +1,260 @@ +#include +#include +#include +#include +#include +#include +#include + +#define PI 3.14159265358979323846 + +namespace py = pybind11; + +std::array matvec(double a11, double a12, double a13, double a21, + double a22, double a23, double a31, double a32, + double a33, const std::array& vec) { + return {a11 * vec[0] + a12 * vec[1] + a13 * vec[2], + a21 * vec[0] + a22 * vec[1] + a23 * vec[2], + a31 * vec[0] + a32 * vec[1] + a33 * vec[2]}; +} + +std::array cross(const std::array& a, + const std::array& b) { + return {a[1] * b[2] - a[2] * b[1], a[2] * b[0] - a[0] * b[2], + a[0] * b[1] - a[1] * b[0]}; +} + +struct FishGridHelper { + FishGridHelper(py::list positions, py::list velocities, py::list directors, + py::list omegas, py::size_t max_surface_grid_density, + py::array_t widths, py::array_t heights) + : num_fish_(positions.size()), + max_surface_grid_density_(max_surface_grid_density) { + /* We assume all fish have the same initial geometry */ + for (const auto& pos : positions) + positions_.emplace_back(pos.cast>().unchecked<2>()); + for (const auto& vel : velocities) + velocities_.emplace_back(vel.cast>().unchecked<2>()); + for (const auto& dir : directors) + directors_.emplace_back(dir.cast>().unchecked<3>()); + for (const auto& omega : omegas) + omegas_.emplace_back(omega.cast>().unchecked<2>()); + num_elements_ = omegas_[0].shape(1); + + /* Compute number of lagrangian nodes */ + auto widths_unchecked = widths.unchecked<1>(); + auto heights_unchecked = heights.unchecked<1>(); + std::vector circumference(num_elements_); + double max_circumference = 0.0; + for (unsigned i = 0; i < num_elements_; ++i) { + circumference[i] = + PI * (3 * (widths_unchecked[i] + heights_unchecked[i]) - + std::sqrt((3 * widths_unchecked[i] + heights_unchecked[i]) * + (3 * heights_unchecked[i] + widths_unchecked[i]))); + max_circumference = std::max(max_circumference, circumference[i]); + } + max_grid_size_ = max_circumference / max_surface_grid_density_; + + std::vector surface_grid_density(num_elements_, 0); + start_index_.reserve(num_elements_ + 1); + start_index_.push_back(0); + for (unsigned i = 0; i < num_elements_; ++i) { + surface_grid_density[i] = std::round( + circumference[i] * max_surface_grid_density_ / max_circumference); + // If there are less than 1 point then set it equal to 1 since we will + // place it on the element center. + if (surface_grid_density[i] < 3) + surface_grid_density[i] = 1; + + num_lag_nodes_ += surface_grid_density[i]; + start_index_.push_back(start_index_.back() + surface_grid_density[i]); + } + + /* Compute the grid positions in the local frame. Note that + * "heights"/"semi_minor_axis" is aligned with director 1, while + * "widths"/"semi_major_axis" is aligned with director 2 */ + local_frame_surface_points_.reserve(num_lag_nodes_); + for (unsigned i = 0; i < num_elements_; ++i) { + // If the element has only one grid point, assign it to element center + if (surface_grid_density[i] == 1) { + local_frame_surface_points_.push_back({0, 0, 0}); + continue; + } + + // Otherwise linspace distribution (no end point) + double dt = 2 * PI / surface_grid_density[i]; + for (unsigned j = 0; j < surface_grid_density[i]; ++j) { + local_frame_surface_points_.push_back( + {heights_unchecked[i] * std::cos(j * dt), + widths_unchecked[i] * std::sin(j * dt), 0.0}); + } + } + + moment_arms_.reserve(num_fish_); + for (unsigned fi = 0; fi < num_fish_; ++fi) { + moment_arms_.push_back({num_lag_nodes_, {0, 0, 0}}); + } + } + + py::size_t get_num_lag_nodes() const { return num_fish_ * num_lag_nodes_; } + + void set_element_length(double length) { + max_grid_size_ = std::max(max_grid_size_, length); + } + + void compute_lag_grid_position_field(py::array_t lag_positions) { + auto lag_pos = lag_positions.mutable_unchecked<2>(); + compute_moment_arms_(); + for (unsigned fi = 0; fi < num_fish_; ++fi) { // for each fish + unsigned start = fi * num_lag_nodes_; + for (unsigned i = 0; i < num_elements_; ++i) { // for each element + std::array elem_pos = { + 0.5 * (positions_[fi](0, i) + positions_[fi](0, i + 1)), + 0.5 * (positions_[fi](1, i) + positions_[fi](1, i + 1)), + 0.5 * (positions_[fi](2, i) + positions_[fi](2, i + 1))}; + for (unsigned j = start_index_[i]; j < start_index_[i + 1]; ++j) { + // for each lagrangian grid point + lag_pos(0, start + j) = elem_pos[0] + moment_arms_[fi][j][0]; + lag_pos(1, start + j) = elem_pos[1] + moment_arms_[fi][j][1]; + lag_pos(2, start + j) = elem_pos[2] + moment_arms_[fi][j][2]; + } // j + } // i + } // fi + } + + void compute_lag_grid_velocity_field(py::array_t lag_velocities) { + auto lag_vel = lag_velocities.mutable_unchecked<2>(); + for (unsigned fi = 0; fi < num_fish_; ++fi) { // for each fish + unsigned start = fi * num_lag_nodes_; + for (unsigned i = 0; i < num_elements_; ++i) { // for each element + std::array elem_vel = { + 0.5 * (velocities_[fi](0, i) + velocities_[fi](0, i + 1)), + 0.5 * (velocities_[fi](1, i) + velocities_[fi](1, i + 1)), + 0.5 * (velocities_[fi](2, i) + velocities_[fi](2, i + 1))}; + std::array elem_glob_omega = + matvec(directors_[fi](0, 0, i), directors_[fi](1, 0, i), + directors_[fi](2, 0, i), directors_[fi](0, 1, i), + directors_[fi](1, 1, i), directors_[fi](2, 1, i), + directors_[fi](0, 2, i), directors_[fi](1, 2, i), + directors_[fi](2, 2, i), + {omegas_[fi](0, i), omegas_[fi](1, i), omegas_[fi](2, i)}); + for (unsigned j = start_index_[i]; j < start_index_[i + 1]; ++j) { + // for each lagrangian grid point + auto local_vel = cross(elem_glob_omega, moment_arms_[fi][j]); + lag_vel(0, start + j) = elem_vel[0] + local_vel[0]; + lag_vel(1, start + j) = elem_vel[1] + local_vel[1]; + lag_vel(2, start + j) = elem_vel[2] + local_vel[2]; + } // j + } // i + } // fi + } + + void transfer_forcing_from_grid_to_body( + py::array_t body_flow_forces, + py::array_t body_flow_torques, + py::array_t lag_grid_forcing_field) { + // Assume body_flow_forces and body_flow_torques are zeroed out + auto body_force = body_flow_forces.mutable_unchecked<2>(); + auto body_torque = body_flow_torques.mutable_unchecked<2>(); + auto lag_forcing = lag_grid_forcing_field.unchecked<2>(); + for (unsigned fi = 0; fi < num_fish_; ++fi) { // for each fish + unsigned start = fi * num_lag_nodes_; + unsigned start_elem = fi * num_elements_; + unsigned start_node = fi * (num_elements_ + 1); + + for (unsigned i = 0; i < num_elements_; ++i) { // for each element + unsigned elem_index = start_elem + i; + unsigned node_index = start_node + i; + + std::array body_force_on_elem{0.0, 0.0, 0.0}; + std::array body_torque_on_elem{0.0, 0.0, 0.0}; + for (unsigned j = start_index_[i]; j < start_index_[i + 1]; ++j) { + // for each lagrangian grid point + unsigned lag_index = start + j; + body_force_on_elem[0] += lag_forcing(0, lag_index); + body_force_on_elem[1] += lag_forcing(1, lag_index); + body_force_on_elem[2] += lag_forcing(2, lag_index); + + auto local_torque = + cross(moment_arms_[fi][j], + {-lag_forcing(0, lag_index), -lag_forcing(1, lag_index), + -lag_forcing(2, lag_index)}); + body_torque_on_elem[0] += local_torque[0]; + body_torque_on_elem[1] += local_torque[1]; + body_torque_on_elem[2] += local_torque[2]; + } // j + + body_force(0, node_index) -= 0.5 * body_force_on_elem[0]; + body_force(1, node_index) -= 0.5 * body_force_on_elem[1]; + body_force(2, node_index) -= 0.5 * body_force_on_elem[2]; + + body_force(0, node_index + 1) -= 0.5 * body_force_on_elem[0]; + body_force(1, node_index + 1) -= 0.5 * body_force_on_elem[1]; + body_force(2, node_index + 1) -= 0.5 * body_force_on_elem[2]; + + body_torque_on_elem = + matvec(directors_[fi](0, 0, i), directors_[fi](0, 1, i), + directors_[fi](0, 2, i), directors_[fi](1, 0, i), + directors_[fi](1, 1, i), directors_[fi](1, 2, i), + directors_[fi](2, 0, i), directors_[fi](2, 1, i), + directors_[fi](2, 2, i), body_torque_on_elem); + body_torque(0, elem_index) = body_torque_on_elem[0]; + body_torque(1, elem_index) = body_torque_on_elem[1]; + body_torque(2, elem_index) = body_torque_on_elem[2]; + } // i + } // fi + } + + double get_maximum_lagrangian_grid_spacing() { + // for (unsigned fi = 0; fi < num_fish_; ++fi) { + // auto lengths_unchecked = + // lengths[fi].cast>().unchecked<1>(); + // for (unsigned i = 0; i < num_elements_; ++i) { + // max_grid_size_ = std::max(max_grid_size_, lengths_unchecked[i]); + // } + // } + return max_grid_size_; + } + + private: + unsigned num_fish_, num_elements_, max_surface_grid_density_, num_lag_nodes_; + double max_grid_size_; + std::vector> positions_; + std::vector> velocities_; + std::vector> directors_; + std::vector> omegas_; + std::vector> local_frame_surface_points_; + std::vector>> moment_arms_; + std::vector start_index_; + + void compute_moment_arms_() { + for (unsigned fi = 0; fi < num_fish_; ++fi) { + for (unsigned i = 0; i < num_elements_; ++i) { + for (unsigned j = start_index_[i]; j < start_index_[i + 1]; ++j) { + moment_arms_[fi][j] = + matvec(directors_[fi](0, 0, i), directors_[fi](1, 0, i), + directors_[fi](2, 0, i), directors_[fi](0, 1, i), + directors_[fi](1, 1, i), directors_[fi](2, 1, i), + directors_[fi](0, 2, i), directors_[fi](1, 2, i), + directors_[fi](2, 2, i), local_frame_surface_points_[j]); + } // j + } // i + } // fi + } +}; + +PYBIND11_MODULE(_fish_grid_helper, m) { + py::class_(m, "FishGridHelper") + .def(py::init, py::array_t>()) + .def("get_num_lag_nodes", &FishGridHelper::get_num_lag_nodes) + .def("set_element_length", &FishGridHelper::set_element_length) + .def("compute_lag_grid_position_field", + &FishGridHelper::compute_lag_grid_position_field) + .def("compute_lag_grid_velocity_field", + &FishGridHelper::compute_lag_grid_velocity_field) + .def("transfer_forcing_from_grid_to_body", + &FishGridHelper::transfer_forcing_from_grid_to_body) + .def("get_maximum_lagrangian_grid_spacing", + &FishGridHelper::get_maximum_lagrangian_grid_spacing); +} \ No newline at end of file diff --git a/examples/fish_swimming/utils/fish_grid.py b/examples/fish_swimming/utils/fish_grid.py new file mode 100644 index 00000000..63bdbbd4 --- /dev/null +++ b/examples/fish_swimming/utils/fish_grid.py @@ -0,0 +1,364 @@ +import numpy as np + +from elastica._linalg import _batch_cross, _batch_matvec, _batch_matrix_transpose +from sopht.simulator.immersed_body import ImmersedBodyForcingGrid + +from elasticapp.Systems.FishModel._PyFishModel import Fish + +from ._fish_grid_helper import FishGridHelper + + +class FishSurfaceForcingGrid(ImmersedBodyForcingGrid): + """Immersed boundary forcing grid for a single fish. + + This is a close replication of the fish swimming example in sopht. + Implementation is purely Pythonic. + """ + + def __init__( + self, + grid_dim: int, + cosserat_rod: Fish, + surface_grid_density_for_largest_element: int, + ) -> None: + if grid_dim != 3: + msg = ( + "Invalid grid dimensions. Cosserat rod surface forcing grid is only " + "defined for grid_dim=3" + ) + raise ValueError(msg) + + self.position_collection = np.asarray(cosserat_rod.position_collection) + self.velocity_collection = np.asarray(cosserat_rod.velocity_collection) + self.omega_collection = np.asarray(cosserat_rod.omega_collection) + self.director_collection = np.asarray(cosserat_rod.director_collection) + self.dilatation = np.asarray(cosserat_rod.dilatation) + self.lengths = np.asarray(cosserat_rod.lengths) + + self.height = np.asarray(cosserat_rod.semi_major_axis) + + # Surface grid density at the arm maximum radius + self.width = np.asarray(cosserat_rod.semi_minor_axis) + self.surface_grid_density_for_largest_element = surface_grid_density_for_largest_element + + # area = 4*np.pi * self.width * self.height + # self.surface_grid_points = np.rint(area / np.max(area) * surface_grid_density_for_largest_element).astype(int) + # Srinivasa Ramanujan approximation for ellipse circumference + # https://en.wikipedia.org/wiki/Ellipse + self.circumference = np.pi * ( + 3 * (self.width + self.height) + - np.sqrt((3 * self.width + self.height) * (self.width + 3 * self.height)) + ) + self.surface_grid_points = np.rint( + self.circumference + / np.max(self.circumference) + * surface_grid_density_for_largest_element + ).astype(int) + + # If there are less than 1 point then set it equal to 1 since we will place + # it on the element center. + self.surface_grid_points[np.where(self.surface_grid_points < 3)[0]] = 1 + num_lag_nodes = self.surface_grid_points.sum() + super().__init__(grid_dim, num_lag_nodes) + self.n_elements = len(self.width) + + self.surface_point_rotation_angle_list = [] + for i in range(self.n_elements): + if self.surface_grid_points[i] > 1: + # If there are more than one point on the surface then compute + # the angle of these points. Surface points are on the local frame + self.surface_point_rotation_angle_list.append( + np.linspace(0, 2 * np.pi, self.surface_grid_points[i], endpoint=False) + ) + else: + # If there is only one point, then that point is on the element + # center so pass empty array. + self.surface_point_rotation_angle_list.append(np.array([])) + + # Since lag grid points are on the surface, for each node we need to compute moment arm. + self.moment_arm = np.zeros_like(self.position_field) + + # Depending on the rod taper number of surface points can vary between elements, + # so we need start_idx and end_idx to slice grid points according to element they belong. + self.start_idx = np.zeros((self.n_elements), dtype=int) + self.end_idx = np.zeros((self.n_elements), dtype=int) + + start_idx_temp = 0 + end_idx_temp = 0 + for i in range(self.n_elements): + self.start_idx[i] = start_idx_temp + start_idx_temp += self.surface_grid_points[i] + end_idx_temp += self.surface_grid_points[i] + self.end_idx[i] = end_idx_temp + + self.local_frame_surface_points = np.zeros_like(self.position_field) + + # Compute surface(grid) point positions on local frame for each element. + # If there is one grid point then that point is on the element center + # so just pass 0.0 + for i, surface_point_rotation_angle in enumerate(self.surface_point_rotation_angle_list): + if surface_point_rotation_angle.size == 0: + # This is true if there is only one point for element, and it is on element center. + self.local_frame_surface_points[:, self.start_idx[i] : self.end_idx[i]] = 0.0 + else: + self.local_frame_surface_points[0, self.start_idx[i] : self.end_idx[i]] = ( + self.height[i] * np.cos(surface_point_rotation_angle) + ) + self.local_frame_surface_points[1, self.start_idx[i] : self.end_idx[i]] = ( + self.width[i] * np.sin(surface_point_rotation_angle) + ) + + # some caching stuff + self.rod_director_collection_transpose = np.zeros_like(self.director_collection) + self.rod_element_position = np.zeros((self.grid_dim, self.n_elements)) + self.rod_element_velocity = np.zeros_like(self.rod_element_position) + self.rod_element_global_frame_omega = np.zeros_like(self.rod_element_position) + + self.grid_point_director_transpose = np.zeros((3, 3, self.num_lag_nodes)) + self.grid_point_radius = np.zeros((self.num_lag_nodes)) + self.grid_point_dilatation = np.zeros((self.num_lag_nodes)) + self.grid_point_omega = np.zeros_like(self.position_field) + self.lag_grid_torque_field = np.zeros_like(self.position_field) + + # to ensure position/velocity are consistent during initialisation + self.compute_lag_grid_position_field() + self.compute_lag_grid_velocity_field() + + def compute_lag_grid_position_field(self) -> None: + """Computes location of forcing grid for the Cosserat rod""" + + self.rod_element_position[...] = 0.5 * ( + self.position_collection[..., 1:] + self.position_collection[..., :-1] + ) + + # Cache rod director collection transpose since it will be used to + # compute velocity field. + self.rod_director_collection_transpose[...] = _batch_matrix_transpose( + self.director_collection + ) + + # Broadcast rod properties to the grid points. + for i in range(self.n_elements): + self.grid_point_director_transpose[:, :, self.start_idx[i] : self.end_idx[i]] = ( + self.rod_director_collection_transpose[:, :, i : i + 1] + ) + self.grid_point_dilatation[self.start_idx[i] : self.end_idx[i]] = self.dilatation[i] + self.position_field[:, self.start_idx[i] : self.end_idx[i]] = self.rod_element_position[ + :, i : i + 1 + ] + + # Compute the moment arm or distance from the element center + # for each grid point. + self.moment_arm[:] = _batch_matvec( + self.grid_point_director_transpose, self.local_frame_surface_points + ) + + # Surface positions are moment_arm + element center position + self.position_field += self.moment_arm + + def compute_lag_grid_velocity_field(self) -> None: + """Computes velocity of forcing grid points for the Cosserat rod""" + + # Element velocity + self.rod_element_velocity[...] = 0.5 * ( + self.velocity_collection[:, 1:] + self.velocity_collection[:, :-1] + ) + # Element angular velocity + self.rod_element_global_frame_omega[...] = _batch_matvec( + self.rod_director_collection_transpose, + self.omega_collection, + ) + + # Broadcast rod properties to the grid points. + for i in range(self.n_elements): + self.grid_point_omega[:, self.start_idx[i] : self.end_idx[i]] = ( + self.rod_element_global_frame_omega[:, i : i + 1] + ) + self.velocity_field[:, self.start_idx[i] : self.end_idx[i]] = self.rod_element_velocity[ + :, i : i + 1 + ] + + # v_elem + omega X moment_arm + self.velocity_field += _batch_cross(self.grid_point_omega, self.moment_arm) + + def transfer_forcing_from_grid_to_body( + self, + body_flow_forces: np.ndarray, + body_flow_torques: np.ndarray, + lag_grid_forcing_field: np.ndarray, + ) -> None: + """Transfer forcing from lagrangian forcing grid to the cosserat rod""" + body_flow_forces[...] = 0.0 + body_flow_torques[...] = 0.0 + + # negative sign due to Newtons third law + for i in range(self.n_elements): + body_forces_on_elems = np.sum( + lag_grid_forcing_field[:, self.start_idx[i] : self.end_idx[i]], axis=1 + ) + body_flow_forces[:, i] -= 0.5 * body_forces_on_elems + body_flow_forces[:, i + 1] -= 0.5 * body_forces_on_elems + + # negative sign due to Newtons third law + # torque generated by all lagrangian points are + self.lag_grid_torque_field[...] = _batch_cross(self.moment_arm, -lag_grid_forcing_field) + + # Update body torques + # convert global to local frame + for i in range(self.n_elements): + body_flow_torques[:, i] = self.director_collection[:, :, i] @ np.sum( + self.lag_grid_torque_field[:, self.start_idx[i] : self.end_idx[i]], + axis=1, + ) + + def get_maximum_lagrangian_grid_spacing(self) -> float: + """Get the maximum Lagrangian grid spacing""" + return np.amax( + [ + self.lengths, + self.circumference / self.surface_grid_density_for_largest_element, + ] + ) + + +class MultiFishSurfaceForcingGrid(ImmersedBodyForcingGrid): + """Immersed boundary forcing grid for multiple fish. + + The implementation is naive: each fish has its own FishSurfaceForcingGrid and + MultiFishSurfaceForcingGrid simply concatenates the grids of each fish into a + single chunk. + """ + + def __init__( + self, + grid_dim: int, + fish_list: list[Fish], + surface_grid_density_for_largest_element: int, + ) -> None: + self.num_fish = len(fish_list) + if self.num_fish == 0: + msg = "No fish provided to the MultiFishSurfaceForcingGrid." + raise ValueError(msg) + + self.sub_grids = [ + FishSurfaceForcingGrid( + grid_dim=grid_dim, + cosserat_rod=fish, + surface_grid_density_for_largest_element=surface_grid_density_for_largest_element, + ) + for fish in fish_list + ] + + num_lag_nodes = 0 + for sub_grid in self.sub_grids: + num_lag_nodes += sub_grid.num_lag_nodes + + super().__init__(grid_dim, num_lag_nodes) + + start = 0 + + for sub_grid in self.sub_grids: + end = start + sub_grid.num_lag_nodes + + self.position_field[:, start:end] = sub_grid.position_field + self.velocity_field[:, start:end] = sub_grid.velocity_field + + del sub_grid.position_field + del sub_grid.velocity_field + + sub_grid.position_field = self.position_field[:, start:end].view() + sub_grid.velocity_field = self.velocity_field[:, start:end].view() + + start = end + + def compute_lag_grid_position_field(self) -> None: + for sub_grid in self.sub_grids: + sub_grid.compute_lag_grid_position_field() + + def compute_lag_grid_velocity_field(self) -> None: + for sub_grid in self.sub_grids: + sub_grid.compute_lag_grid_velocity_field() + + def transfer_forcing_from_grid_to_body( + self, + body_flow_forces: np.ndarray, + body_flow_torques: np.ndarray, + lag_grid_forcing_field: np.ndarray, + ) -> None: + start_elem = start_node = start_lag = 0 + for sub_grid in self.sub_grids: + end_elem = start_elem + sub_grid.n_elements + end_node = start_node + sub_grid.n_elements + 1 + end_lag = start_lag + sub_grid.num_lag_nodes + sub_grid.transfer_forcing_from_grid_to_body( + body_flow_forces=body_flow_forces[:, start_node:end_node].view(), + body_flow_torques=body_flow_torques[:, start_elem:end_elem].view(), + lag_grid_forcing_field=lag_grid_forcing_field[:, start_lag:end_lag].view(), + ) + start_elem = end_elem + start_node = end_node + start_lag = end_lag + + def get_maximum_lagrangian_grid_spacing(self) -> float: + max_spacing = 0.0 + for sub_grid in self.sub_grids: + max_spacing = max(max_spacing, sub_grid.get_maximum_lagrangian_grid_spacing()) + + return max_spacing + + +class MultiFishForcingGrid(ImmersedBodyForcingGrid): + """Immersed boundary forcing grid for multiple fish. + + This implementation uses C++ helper and is much more optimized than + MultiFishSurfaceForcingGrid. + """ + def __init__( + self, + grid_dim: int, + fish_list: list[Fish], + surface_grid_density_for_largest_element: int, + ) -> None: + position_list = [np.asarray(fish.position_collection) for fish in fish_list] + velocity_list = [np.asarray(fish.velocity_collection) for fish in fish_list] + director_list = [np.asarray(fish.director_collection) for fish in fish_list] + omega_list = [np.asarray(fish.omega_collection) for fish in fish_list] + widths = np.asarray(fish_list[0].semi_minor_axis) + heights = np.asarray(fish_list[0].semi_major_axis) + + self.helper = FishGridHelper( + position_list, + velocity_list, + director_list, + omega_list, + surface_grid_density_for_largest_element, + widths, + heights, + ) + self.helper.set_element_length(np.asarray(fish_list[0].lengths)[0]) + + super().__init__( + grid_dim=grid_dim, + num_lag_nodes=int(self.helper.get_num_lag_nodes()), + ) + self.compute_lag_grid_position_field() + self.compute_lag_grid_velocity_field() + + def compute_lag_grid_position_field(self) -> None: + self.helper.compute_lag_grid_position_field(self.position_field) + + def compute_lag_grid_velocity_field(self) -> None: + self.helper.compute_lag_grid_velocity_field(self.velocity_field) + + def transfer_forcing_from_grid_to_body( + self, body_flow_forces, body_flow_torques, lag_grid_forcing_field + ): + body_flow_forces[...] = 0.0 + body_flow_torques[...] = 0.0 + + self.helper.transfer_forcing_from_grid_to_body( + body_flow_forces, body_flow_torques, lag_grid_forcing_field + ) + + def get_maximum_lagrangian_grid_spacing(self) -> float: + return self.helper.get_maximum_lagrangian_grid_spacing() diff --git a/examples/fish_swimming/utils/fish_io.py b/examples/fish_swimming/utils/fish_io.py new file mode 100644 index 00000000..39d56306 --- /dev/null +++ b/examples/fish_swimming/utils/fish_io.py @@ -0,0 +1,364 @@ +from warnings import warn + +import numpy as np +import h5py + +from elasticapp.Systems.FishModel._PyFishModel import Fish + + +class ImageDataIO: + def __init__( + self, + dx: float, + grid_size: tuple[int, int, int], + ) -> None: + self.grid_size = grid_size + self.origin = (0.5 * dx,) * 3 + self.spacing = (dx,) * 3 + + def write_to_file(self, filename: str, time: float | np.float64, **kwargs) -> None: + with h5py.File(filename, "w") as f: + image_data_group = f.create_group(name="VTKHDF") + image_data_group.attrs["Version"] = (2, 2) + image_data_group.attrs["Type"] = "ImageData".encode("ascii") + image_data_group.attrs["WholeExtent"] = ( + 0, + self.grid_size[2] - 1, + 0, + self.grid_size[1] - 1, + 0, + self.grid_size[0] - 1, + ) + image_data_group.attrs["Direction"] = ( + 1.0, + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + 0.0, + 1.0, + ) + image_data_group.attrs["Origin"] = self.origin + image_data_group.attrs["Spacing"] = self.spacing + + point_data_group = image_data_group.create_group(name="PointData") + # cell_data_group = image_data_group.create_group(name="CellData") + field_data_group = image_data_group.create_group(name="FieldData") + + field_data_group.create_dataset( + name="Time", + shape=(1,), + dtype=np.float64, + data=np.array([time], dtype=np.float64), + ) + + for key, value in kwargs.items(): + if not isinstance(value, np.ndarray): + warn(f"Value for {key} is not a numpy array. Skipped.") + continue + + if len(value.shape) != 3 and len(value.shape) != 4: + warn(f"Value for {key} is not a 3D or 4D array. Skipped.") + continue + + if len(value.shape) == 3: + if value.shape != self.grid_size: + warn( + f"Value for {key} has shape {value.shape} " + f"but expected {self.grid_size}. Skipped." + ) + continue + + point_data_group.create_dataset( + name=key, + shape=self.grid_size, + dtype=np.float64, + data=value.copy(), + ) + + if len(value.shape) == 4: + if value.shape != (3,) + self.grid_size: + warn( + f"Value for {key} has shape {value.shape} " + f"but expected {(3,) + self.grid_size}. Skipped." + ) + continue + + point_data_group.create_dataset( + name=key, + shape=(*self.grid_size, 3), + dtype=np.float64, + data=np.transpose(value, (1, 2, 3, 0)).copy(), + ) + + def read_from_file(self, filename: str, **kwargs) -> np.float64: + time = np.float64(-1.0) + with h5py.File(filename, "r", driver="mpio", comm=self.mpi_comm) as f: + time = f["/VTKHDF/FieldData/Time"][0] + for field_name, field_view in kwargs.items(): + path = "/VTKHDF/PointData/" + field_name + try: + dataset = f[path] + except KeyError: + msg = f"Dataset '{field_name}' not found in PointData of file {filename}" + raise KeyError(msg) + + if dataset.shape != self.global_grid_size and dataset.shape != ( + *self.global_grid_size, + 3, + ): + msg = ( + f"Dataset '{field_name}' has shape {dataset.shape} " + f"but expected {self.global_grid_size} or (*{self.global_grid_size}, 3)." + ) + raise ValueError(msg) + + if len(dataset.shape) == 3: + field_view[...] = dataset + + elif len(dataset.shape) == 4: + field_view[...] = np.transpose(dataset, (3, 0, 1, 2)) + + return time + + +class FishIO: + def __init__( + self, + fish_list: list[Fish], + virtual_fish_list: list[Fish], + ) -> None: + self.num_fish = len(fish_list) + if self.num_fish <= 0: + msg = "fish_list must have at least one fish." + raise ValueError(msg) + if self.num_fish != len(virtual_fish_list): + msg = "virtual_fish_list must have the same length as fish_list." + raise ValueError(msg) + + self.n_elem = int(fish_list[0].n_elems[0]) + for _n_elem in [int(fish.n_elems[0]) for fish in fish_list + virtual_fish_list]: + if self.n_elem != _n_elem: + msg = "All fish must have the same number of elements." + raise ValueError(msg) + + self.fish_dicts = [] + for i in range(self.num_fish): + self.fish_dicts.append( + { + "name": f"fish_{i}", + "semi_major_axis": np.asarray(fish_list[i].semi_major_axis), + "semi_minor_axis": np.asarray(fish_list[i].semi_minor_axis), + "mass": np.asarray(fish_list[i].mass), + "position": np.asarray(fish_list[i].position_collection), + "director": np.asarray(fish_list[i].director_collection), + "velocity": np.asarray(fish_list[i].velocity_collection), + "angular_velocity": np.asarray(fish_list[i].omega_collection), + "v_position": np.asarray(virtual_fish_list[i].position_collection), + "v_director": np.asarray(virtual_fish_list[i].director_collection), + "v_velocity": np.asarray(virtual_fish_list[i].velocity_collection), + "v_angular_velocity": np.asarray(virtual_fish_list[i].omega_collection), + } + ) + + def write_to_file(self, filename: str, time: float | np.float64) -> None: + with h5py.File(filename, "w") as f: + root = f.create_group("root") + root.create_dataset( + name="Time", + shape=(1,), + dtype=np.float64, + data=np.array([time], dtype=np.float64), + ) + + for fish_dict in self.fish_dicts: + group = root.create_group(fish_dict["name"]) + for field_name, field_array in fish_dict.items(): + if field_name == "name": + continue + group.create_dataset( + name=field_name, + shape=field_array.shape, + dtype=np.float64, + data=field_array.copy(), + ) + + def read_from_file(self, filename: str) -> np.float64: + time = np.float64(-1.0) + with h5py.File(filename, "r") as f: + time = f["/root/Time"][0] + for fish_dict in self.fish_dicts: + path = "/root/" + fish_dict["name"] + try: + group = f[path] + except KeyError: + msg = f"Group '{fish_dict['name']}' not found in file {filename} at {path}" + raise KeyError(msg) + + for key, value in fish_dict.items(): + if key != "name": + try: + value[...] = group[key] + except KeyError: + msg = f"Dataset '{key}' not found in group in file {filename} at {path}/{key}" + raise KeyError(msg) + + return time + + +class GridIO: + def __init__( + self, + num_fish: int, + num_lag_points_per_fish: int, + ) -> None: + self.num_fish = num_fish + + self.connectivities = ["Vertices", "Lines", "Polygons", "Strips"] + self.num_points = num_fish * num_lag_points_per_fish + self.num_cells = num_fish * (num_lag_points_per_fish - 1) + + self.num_line_conn = 2 * self.num_cells + self.line_connectivity = np.empty((self.num_line_conn,), dtype=np.int32) + + idx_start = 0 + pos_start = 0 + for _ in range(num_fish): + for j in range(num_lag_points_per_fish - 1): + self.line_connectivity[idx_start + 2 * j] = pos_start + j + self.line_connectivity[idx_start + 2 * j + 1] = pos_start + j + 1 + + idx_start += 2 * (num_lag_points_per_fish - 1) + pos_start += num_lag_points_per_fish + + self.line_offsets = np.arange(0, self.num_line_conn + 1, 2, dtype=np.int32) + + self.filter_kwargs = { + "compression": "gzip", + "compression_opts": 9, + } + + def write_to_file( + self, filename: str, time: float | np.float64, position: np.ndarray, **kwargs + ) -> None: + if position.shape != (3, self.num_points): + msg = f"Position array has shape {position.shape} but expected (3, {self.num_points})." + raise ValueError(msg) + + with h5py.File(filename, "w") as f: + root = f.create_group("VTKHDF") + root.attrs["Version"] = (2, 2) + ascii_type = "PolyData".encode("ascii") + root.attrs.create("Type", ascii_type, dtype=h5py.string_dtype("ascii", len(ascii_type))) + + root.create_dataset( + "NumberOfPoints", + data=np.array([self.num_points], dtype=np.int32), + ) + root.create_dataset( + "Points", + data=position.T, + dtype=np.float64, + **self.filter_kwargs, + ) + + for conn_type in self.connectivities: + conn_group = root.create_group(conn_type) + if conn_type == "Lines": + conn_group.create_dataset( + "NumberOfConnectivityIds", + data=np.array([self.num_line_conn], dtype=np.int32), + ) + conn_group.create_dataset( + "NumberOfCells", + data=np.array([self.num_cells], dtype=np.int32), + ) + conn_group.create_dataset( + "Connectivity", + data=self.line_connectivity, + dtype=np.int32, + **self.filter_kwargs, + ) + conn_group.create_dataset( + "Offsets", + data=self.line_offsets, + dtype=np.int32, + **self.filter_kwargs, + ) + + else: + conn_group.create_dataset( + "NumberOfConnectivityIds", + data=np.array([0], dtype=np.int32), + ) + conn_group.create_dataset( + "NumberOfCells", + data=np.array([0], dtype=np.int32), + ) + conn_group.create_dataset( + "Connectivity", + data=np.array([], dtype=np.int32), + ) + conn_group.create_dataset( + "Offsets", + data=np.array([0], dtype=np.int32), + ) + + # Point data + point_data = root.create_group("PointData") + for key, value in kwargs.items(): + if not isinstance(value, np.ndarray): + msg = f"Value for {key} is not a numpy array." + raise ValueError(msg) + if value.shape != (3, self.num_points): + msg = f"Value for {key} has shape {value.shape} but expected (3, {self.num_points})." + raise ValueError(msg) + point_data.create_dataset( + key, + data=value.T, + dtype=np.float64, + **self.filter_kwargs, + ) + + # Field data + field_data = root.create_group("FieldData") + field_data.create_dataset( + "time", + data=np.array([time], dtype=np.float64), + ) + + def read_from_file(self, filename: str, **kwargs) -> np.float64: + with h5py.File(filename, "r") as file: + root = file["VTKHDF"] + time = root["FieldData"]["time"][0] + + for key, value in kwargs.items(): + if key == "position": + dataset = root["Points"] + if dataset.shape != (self.num_points, 3): + msg = f"Dataset 'Points' has shape {dataset.shape} but expected ({self.num_points}, 3)." + raise ValueError(msg) + + if value.shape != (3, self.num_points): + msg = f"Value for 'position' has shape {value.shape} but expected (3, {self.num_points})." + raise ValueError(msg) + value[...] = dataset.T + + elif key not in root["PointData"].keys(): + msg = f"Dataset '{key}' not found in PointData of file {filename}." + raise KeyError(msg) + + elif value.shape != (3, self.num_points): + msg = f"Value for {key} has shape {value.shape} but expected (3, {self.num_points})." + raise ValueError(msg) + + else: + dataset = root["PointData"][key] + if dataset.shape != (self.num_points, 3): + msg = f"Dataset '{key}' has shape {dataset.shape} but expected ({self.num_points}, 3)." + raise ValueError(msg) + value[...] = dataset.T + + return time diff --git a/examples/fish_swimming/utils/fish_io_mpi.py b/examples/fish_swimming/utils/fish_io_mpi.py new file mode 100644 index 00000000..753cea9b --- /dev/null +++ b/examples/fish_swimming/utils/fish_io_mpi.py @@ -0,0 +1,338 @@ +from warnings import warn + +import numpy as np +import h5py +from mpi4py.MPI import Intracomm + +from elasticapp.Systems.FishModel._PyFishModel import Fish + + +class MPIImageDataIO: + def __init__( + self, + mpi_comm: Intracomm, + dx: float, + global_grid_size: tuple[int, int, int], + local_grid_size: tuple[int, int, int], + ghost_size: int, + ) -> None: + self.mpi_comm = mpi_comm + self.global_grid_size = global_grid_size + self.origin = (0.5 * dx,) * 3 + self.spacing = (dx,) * 3 + + start_idx = np.asarray(self.mpi_comm.coords) * local_grid_size + end_idx = start_idx + local_grid_size + self.slice = tuple(slice(start_idx[i], end_idx[i]) for i in range(3)) + ( + slice(None), + ) + self.no_ghost_slice = (slice(None),) + (slice(ghost_size, -ghost_size),) * 3 + self.local_size_with_ghost = tuple( + local_grid_size[i] + 2 * ghost_size for i in range(3) + ) + + def write_to_file(self, filename: str, time: float | np.float64, **kwargs) -> None: + with h5py.File(filename, "w", driver="mpio", comm=self.mpi_comm) as f: + image_data_group = f.create_group(name="VTKHDF") + image_data_group.attrs["Version"] = (2, 2) + image_data_group.attrs["Type"] = np.string_("ImageData") + image_data_group.attrs["WholeExtent"] = ( + 0, + self.global_grid_size[2] - 1, + 0, + self.global_grid_size[1] - 1, + 0, + self.global_grid_size[0] - 1, + ) + image_data_group.attrs["Direction"] = ( + 1.0, + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + 0.0, + 1.0, + ) + image_data_group.attrs["Origin"] = self.origin + image_data_group.attrs["Spacing"] = self.spacing + + point_data_group = image_data_group.create_group(name="PointData") + # cell_data_group = image_data_group.create_group(name="CellData") + field_data_group = image_data_group.create_group(name="FieldData") + + time_data = field_data_group.create_dataset( + name="Time", + shape=(1,), + dtype=np.float64, + ) + if self.mpi_comm.Get_rank() == 0: + time_data[0] = np.float64(time) + + for key, value in kwargs.items(): + if not isinstance(value, np.ndarray): + warn(f"Value for {key} is not a numpy array. Skipped.") + continue + + if len(value.shape) != 3 and len(value.shape) != 4: + warn(f"Value for {key} is not a 3D or 4D array. Skipped.") + continue + + if len(value.shape) == 3: + if value.shape != self.local_size_with_ghost: + warn( + f"Value for {key} has shape {value.shape} " + f"but expected {self.local_size_with_ghost}. Skipped." + ) + continue + + data_to_fill = point_data_group.create_dataset( + name=key, + shape=self.global_grid_size, + dtype=np.float64, + ) + data_to_fill[self.slice[:-1]] = value[ + self.no_ghost_slice[1:] + ].copy() + + if len(value.shape) == 4: + if value.shape != (3,) + self.local_size_with_ghost: + warn( + f"Value for {key} has shape {value.shape} " + f"but expected {(3,) + self.local_size_with_ghost}. Skipped." + ) + continue + + data_to_fill = point_data_group.create_dataset( + name=key, + shape=(*self.global_grid_size, 3), + dtype=np.float64, + ) + data_to_fill[self.slice] = np.transpose( + value[self.no_ghost_slice], (1, 2, 3, 0) + ).copy() + + def read_from_file(self, filename: str, **kwargs) -> np.float64: + time = np.float64(-1.0) + with h5py.File(filename, "r", driver="mpio", comm=self.mpi_comm) as f: + try: + time = f["/VTKHDF/FieldData/Time"][0] + except KeyError: + msg = ( + f"Process {self.mpi_comm.Get_rank()}: " + f"Time not found in FiledData of file {filename}" + ) + raise KeyError(msg) + + for field_name, field_view in kwargs.items(): + path = "/VTKHDF/PointData/" + field_name + try: + dataset = f[path] + except KeyError: + msg = ( + f"Process {self.mpi_comm.Get_rank()}: " + f"Dataset '{field_name}' not found in PointData of file {filename}" + ) + raise KeyError(msg) + + if dataset.shape != self.global_grid_size and dataset.shape != ( + *self.global_grid_size, + 3, + ): + msg = ( + f"Process {self.mpi_comm.Get_rank()}: " + f"Dataset '{field_name}' has shape {dataset.shape} " + f"but expected {self.global_grid_size} or (*{self.global_grid_size}, 3)." + ) + raise ValueError(msg) + + if len(dataset.shape) == 3: + field_view[self.no_ghost_slice[1:]] = dataset[self.slice[:-1]] + + elif len(dataset.shape) == 4: + field_view[self.no_ghost_slice] = np.transpose( + dataset[self.slice], (3, 0, 1, 2) + ) + + return time + + +class MPIFishIO: + def __init__( + self, + mpi_comm: Intracomm, + fish_list: list[Fish], + virtual_fish_list: list[Fish], + index_list: list[int], + num_fish_element: int, + num_lag_nodes: int, + global_num_fish: int, + **lag_fields: dict[str, np.ndarray], + ) -> None: + self.local_num_fish = len(fish_list) + if self.local_num_fish != len(index_list): + msg = "index_list must have the same length as fish_list." + raise ValueError(msg) + if self.local_num_fish != len(virtual_fish_list): + msg = "virtual_fish_list must have the same length as fish_list." + raise ValueError(msg) + + for _, lag_field in lag_fields.items(): + if self.local_num_fish > 0 and lag_field.shape != ( + 3, + self.local_num_fish * num_lag_nodes, + ): + msg = ( + f"Incompatible lag field shape: {lag_field.shape}. " + f"Expected shape: (3, {self.local_num_fish * num_lag_nodes})." + ) + raise ValueError(msg) + + self.mpi_comm = mpi_comm + self.local_num_fish_element = num_fish_element + self.global_num_fish = global_num_fish + self.num_lag_nodes = num_lag_nodes + + self.fish_dicts = [] + for i in range(self.local_num_fish): + self.fish_dicts.append( + { + "name": f"fish_{index_list[i]}", + "position": np.asarray(fish_list[i].position_collection), + "director": np.asarray(fish_list[i].director_collection), + "velocity": np.asarray(fish_list[i].velocity_collection), + "angular_velocity": np.asarray(fish_list[i].omega_collection), + "v_position": np.asarray(virtual_fish_list[i].position_collection), + "v_director": np.asarray(virtual_fish_list[i].director_collection), + "v_velocity": np.asarray(virtual_fish_list[i].velocity_collection), + "v_angular_velocity": np.asarray( + virtual_fish_list[i].omega_collection + ), + } + ) + for key, value in lag_fields.items(): + if key in [ + "name", + "position", + "director", + "velocity", + "angular_velocity", + "v_position", + "v_director", + "v_velocity", + "v_angular_velocity", + ]: + warn( + f"Lagrangian field '{key}' is already used in fish_dict. Skipped." + ) + + else: + self.fish_dicts[-1][key] = value[ + :, i * num_lag_nodes : (i + 1) * num_lag_nodes + ].view() + + self.lag_field_names = list(lag_fields.keys()) + + def write_to_file(self, filename: str, time: float | np.float64) -> None: + with h5py.File(filename, "w", driver="mpio", comm=self.mpi_comm) as f: + root = f.create_group("root") + time_data = root.create_dataset( + name="Time", + shape=(1,), + dtype=np.float64, + ) + if self.mpi_comm.Get_rank() == 0: + time_data[0] = np.float64(time) + + for i in range(self.global_num_fish): + group = root.create_group(f"fish_{i}") + group.create_dataset( + name="position", + shape=(3, self.local_num_fish_element + 1), + dtype=np.float64, + ) + group.create_dataset( + name="director", + shape=(3, 3, self.local_num_fish_element), + dtype=np.float64, + ) + group.create_dataset( + name="velocity", + shape=(3, self.local_num_fish_element + 1), + dtype=np.float64, + ) + group.create_dataset( + name="angular_velocity", + shape=(3, self.local_num_fish_element), + dtype=np.float64, + ) + group.create_dataset( + name="v_position", + shape=(3, self.local_num_fish_element + 1), + dtype=np.float64, + ) + group.create_dataset( + name="v_director", + shape=(3, 3, self.local_num_fish_element), + dtype=np.float64, + ) + group.create_dataset( + name="v_velocity", + shape=(3, self.local_num_fish_element + 1), + dtype=np.float64, + ) + group.create_dataset( + name="v_angular_velocity", + shape=(3, self.local_num_fish_element), + dtype=np.float64, + ) + + for lag_field_name in self.lag_field_names: + group.create_dataset( + name=lag_field_name, + shape=(3, self.num_lag_nodes), + dtype=np.float64, + ) + + for fish_dict in self.fish_dicts: + for key, value in fish_dict.items(): + if key != "name": + path = "/root/" + fish_dict["name"] + "/" + key + root[path][...] = value + + def read_from_file(self, filename: str) -> np.float64: + time = np.float64(-1.0) + with h5py.File(filename, "r", driver="mpio", comm=self.mpi_comm) as f: + try: + time = f["/root/Time"][0] + except KeyError: + msg = ( + f"Process {self.mpi_comm.Get_rank()}: " + f"Time not found in file {filename} at /root/Time" + ) + raise KeyError(msg) + + for fish_dict in self.fish_dicts: + path = "/root/" + fish_dict["name"] + try: + group = f[path] + except KeyError: + msg = ( + f"Process {self.mpi_comm.Get_rank()}: " + f"Group '{fish_dict['name']}' not found in file {filename} at {path}" + ) + raise KeyError(msg) + + for key, value in fish_dict.items(): + if key != "name": + try: + value[...] = group[key] + except KeyError: + msg = ( + f"Process {self.mpi_comm.Get_rank()}: " + f"Dataset '{key}' not found in group in file {filename} at {path}/{key}" + ) + raise KeyError(msg) + + return time diff --git a/examples/fish_swimming/utils/path_utils.py b/examples/fish_swimming/utils/path_utils.py new file mode 100644 index 00000000..57221216 --- /dev/null +++ b/examples/fish_swimming/utils/path_utils.py @@ -0,0 +1,23 @@ +from pathlib import Path + + +def get_data_path(data_dir: str | None = None) -> Path: + if data_dir is not None: + path = Path(data_dir).expanduser().resolve() + else: + path = Path.cwd() / "data" + + path.mkdir(parents=True, exist_ok=True) + return path + + +def get_most_recent_data_file(data_dir: str, file_pattern: str) -> str: + data_path = Path(data_dir).expanduser().resolve() + files = sorted(data_path.glob(file_pattern)) + + if not files: + raise FileNotFoundError( + f"No files matching pattern '{file_pattern}' in {data_path}" + ) + + return str(files[-1])