From 76aa17c704aea7d91de666a887135f9367a0c221 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sat, 27 Apr 2024 07:02:10 -0500 Subject: [PATCH 001/134] update: add cli for mypy run and cleanup --- Makefile | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/Makefile b/Makefile index 204926803..bcdfb6fed 100644 --- a/Makefile +++ b/Makefile @@ -61,6 +61,10 @@ autoflake8-format: .PHONY: format-codestyle format-codestyle: black flake8 +.PHONY: mypy +mypy: + poetry run mypy --config-file pyproject.toml elastica + .PHONY: test test: poetry run pytest @@ -81,7 +85,7 @@ formatting: format-codestyle .PHONY: update-dev-deps update-dev-deps: - poetry add -D pytest@latest coverage@latest pytest-html@latest pytest-cov@latest black@latest + poetry add -D mypy@latest pytest@latest coverage@latest pytest-html@latest pytest-cov@latest black@latest #* Cleaning .PHONY: pycache-remove @@ -92,6 +96,10 @@ pycache-remove: dsstore-remove: find . | grep -E ".DS_Store" | xargs rm -rf +.PHONY: mypycache-remove +mypycache-remove: + find . | grep -E ".mypy_cache" | xargs rm -rf + .PHONY: ipynbcheckpoints-remove ipynbcheckpoints-remove: find . | grep -E ".ipynb_checkpoints" | xargs rm -rf @@ -105,7 +113,7 @@ build-remove: rm -rf build/ .PHONY: cleanup -cleanup: pycache-remove dsstore-remove ipynbcheckpoints-remove pytestcache-remove +cleanup: pycache-remove dsstore-remove ipynbcheckpoints-remove pytestcache-remove mypycache-remove all: format-codestyle cleanup test From 6a4dd5f62a239fac2877bbda6ab5993bbed1926c Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sat, 27 Apr 2024 07:10:45 -0500 Subject: [PATCH 002/134] feat: add mypy checker --- .github/workflows/main.yml | 4 +++ .gitignore | 2 ++ poetry.lock | 51 ++++++++++++++++++++++++++++++++++++-- pyproject.toml | 29 ++++++++++++++++++++++ 4 files changed, 84 insertions(+), 2 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 6c32926fd..973081f63 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -66,6 +66,10 @@ jobs: - name: Run tests run: | make test + - name: Typechecking + if: ${{ startsWith(runner.os, 'macOS') & runner.python-version == '3.10'}} + run: | + make mypy report-coverage: # Report coverage from python 3.10 and mac-os. May change later runs-on: ${{ matrix.os }} strategy: diff --git a/.gitignore b/.gitignore index 62305a521..da20c6d3e 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,8 @@ __pycache__/ # C extensions *.so +*.swp + # Distribution / packaging .Python build/ diff --git a/poetry.lock b/poetry.lock index a45d2d19b..29de342cb 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.5.1 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand. [[package]] name = "accessible-pygments" @@ -947,6 +947,53 @@ files = [ {file = "mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba"}, ] +[[package]] +name = "mypy" +version = "1.10.0" +description = "Optional static typing for Python" +optional = false +python-versions = ">=3.8" +files = [ + {file = "mypy-1.10.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:da1cbf08fb3b851ab3b9523a884c232774008267b1f83371ace57f412fe308c2"}, + {file = "mypy-1.10.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:12b6bfc1b1a66095ab413160a6e520e1dc076a28f3e22f7fb25ba3b000b4ef99"}, + {file = "mypy-1.10.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9e36fb078cce9904c7989b9693e41cb9711e0600139ce3970c6ef814b6ebc2b2"}, + {file = "mypy-1.10.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:2b0695d605ddcd3eb2f736cd8b4e388288c21e7de85001e9f85df9187f2b50f9"}, + {file = "mypy-1.10.0-cp310-cp310-win_amd64.whl", hash = "sha256:cd777b780312ddb135bceb9bc8722a73ec95e042f911cc279e2ec3c667076051"}, + {file = "mypy-1.10.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:3be66771aa5c97602f382230165b856c231d1277c511c9a8dd058be4784472e1"}, + {file = "mypy-1.10.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:8b2cbaca148d0754a54d44121b5825ae71868c7592a53b7292eeb0f3fdae95ee"}, + {file = "mypy-1.10.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1ec404a7cbe9fc0e92cb0e67f55ce0c025014e26d33e54d9e506a0f2d07fe5de"}, + {file = "mypy-1.10.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e22e1527dc3d4aa94311d246b59e47f6455b8729f4968765ac1eacf9a4760bc7"}, + {file = "mypy-1.10.0-cp311-cp311-win_amd64.whl", hash = "sha256:a87dbfa85971e8d59c9cc1fcf534efe664d8949e4c0b6b44e8ca548e746a8d53"}, + {file = "mypy-1.10.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:a781f6ad4bab20eef8b65174a57e5203f4be627b46291f4589879bf4e257b97b"}, + {file = "mypy-1.10.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b808e12113505b97d9023b0b5e0c0705a90571c6feefc6f215c1df9381256e30"}, + {file = "mypy-1.10.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8f55583b12156c399dce2df7d16f8a5095291354f1e839c252ec6c0611e86e2e"}, + {file = "mypy-1.10.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:4cf18f9d0efa1b16478c4c129eabec36148032575391095f73cae2e722fcf9d5"}, + {file = "mypy-1.10.0-cp312-cp312-win_amd64.whl", hash = "sha256:bc6ac273b23c6b82da3bb25f4136c4fd42665f17f2cd850771cb600bdd2ebeda"}, + {file = "mypy-1.10.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:9fd50226364cd2737351c79807775136b0abe084433b55b2e29181a4c3c878c0"}, + {file = "mypy-1.10.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:f90cff89eea89273727d8783fef5d4a934be2fdca11b47def50cf5d311aff727"}, + {file = "mypy-1.10.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fcfc70599efde5c67862a07a1aaf50e55bce629ace26bb19dc17cece5dd31ca4"}, + {file = "mypy-1.10.0-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:075cbf81f3e134eadaf247de187bd604748171d6b79736fa9b6c9685b4083061"}, + {file = "mypy-1.10.0-cp38-cp38-win_amd64.whl", hash = "sha256:3f298531bca95ff615b6e9f2fc0333aae27fa48052903a0ac90215021cdcfa4f"}, + {file = "mypy-1.10.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:fa7ef5244615a2523b56c034becde4e9e3f9b034854c93639adb667ec9ec2976"}, + {file = "mypy-1.10.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:3236a4c8f535a0631f85f5fcdffba71c7feeef76a6002fcba7c1a8e57c8be1ec"}, + {file = "mypy-1.10.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4a2b5cdbb5dd35aa08ea9114436e0d79aceb2f38e32c21684dcf8e24e1e92821"}, + {file = "mypy-1.10.0-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:92f93b21c0fe73dc00abf91022234c79d793318b8a96faac147cd579c1671746"}, + {file = "mypy-1.10.0-cp39-cp39-win_amd64.whl", hash = "sha256:28d0e038361b45f099cc086d9dd99c15ff14d0188f44ac883010e172ce86c38a"}, + {file = "mypy-1.10.0-py3-none-any.whl", hash = "sha256:f8c083976eb530019175aabadb60921e73b4f45736760826aa1689dda8208aee"}, + {file = "mypy-1.10.0.tar.gz", hash = "sha256:3d087fcbec056c4ee34974da493a826ce316947485cef3901f511848e687c131"}, +] + +[package.dependencies] +mypy-extensions = ">=1.0.0" +tomli = {version = ">=1.1.0", markers = "python_version < \"3.11\""} +typing-extensions = ">=4.1.0" + +[package.extras] +dmypy = ["psutil (>=4.0)"] +install-types = ["pip"] +mypyc = ["setuptools (>=50)"] +reports = ["lxml"] + [[package]] name = "mypy-extensions" version = "1.0.0" @@ -1967,4 +2014,4 @@ examples = ["cma"] [metadata] lock-version = "2.0" python-versions = ">=3.10,<3.12" -content-hash = "1345805f189df83c1d4f3e244094e03fa602499461ced2eaf15e85d27780c762" +content-hash = "a59aa0afdd2de945a3204b0fa13810185b7a6201b6b3f616b25d8440b9fe5318" diff --git a/pyproject.toml b/pyproject.toml index 7661ef7b1..1f6edb741 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,6 +52,8 @@ myst-parser = {version = "^1.0", optional = true} numpydoc = {version = "^1.3.1", optional = true} docutils = {version = "^0.18", optional = true} cma = {version = "^3.2.2", optional = true} +mypy = "^1.10.0" +mypy-extensions = "^1.0.0" [tool.poetry.dev-dependencies] black = "21.12b0" @@ -104,3 +106,30 @@ exclude = ''' # https://docs.pytest.org/en/6.2.x/customize.html#pyproject-toml # Directories that are not visited by pytest collector: norecursedirs =["hooks", "*.egg", ".eggs", "dist", "build", "docs", ".tox", ".git", "__pycache__"] + +[tool.mypy] +# https://mypy.readthedocs.io/en/latest/config_file.html#using-a-pyproject-toml-file +python_version = 3.10 +pretty = true +show_traceback = true +color_output = true +strict = true + +allow_redefinition = false +check_untyped_defs = true +disallow_any_generics = false +disallow_incomplete_defs = true +ignore_missing_imports = true +implicit_reexport = true +no_implicit_optional = true +show_column_numbers = true +show_error_codes = true +show_error_context = true +strict_equality = true +strict_optional = true +warn_no_return = true +warn_redundant_casts = true +warn_return_any = true +warn_unreachable = true +warn_unused_configs = true +warn_unused_ignores = true From e75ca6b576a77afaf211a32b96504892b83885e2 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sat, 27 Apr 2024 07:15:33 -0500 Subject: [PATCH 003/134] Update main.yml --- .github/workflows/main.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 973081f63..782202d95 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -67,7 +67,7 @@ jobs: run: | make test - name: Typechecking - if: ${{ startsWith(runner.os, 'macOS') & runner.python-version == '3.10'}} + if: ${{ startsWith(runner.os, 'macOS') && runner.python-version == '3.10'}} run: | make mypy report-coverage: # Report coverage from python 3.10 and mac-os. May change later From e2f2e730dd44e41db11a8898f60b33a8eb023264 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sat, 27 Apr 2024 11:43:16 -0500 Subject: [PATCH 004/134] Update main.yml --- .github/workflows/main.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 782202d95..06bd9852c 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -67,7 +67,7 @@ jobs: run: | make test - name: Typechecking - if: ${{ startsWith(runner.os, 'macOS') && runner.python-version == '3.10'}} + if: ${{ startsWith(runner.os, 'macOS') }} run: | make mypy report-coverage: # Report coverage from python 3.10 and mac-os. May change later From f1b71268bceecc20ea790de258e6f0d78e1b270d Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 28 Apr 2024 02:11:06 -0500 Subject: [PATCH 005/134] update mypy setup --- pyproject.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 1f6edb741..a2c7bf844 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -119,6 +119,8 @@ allow_redefinition = false check_untyped_defs = true disallow_any_generics = false disallow_incomplete_defs = true +disallow_untyped_calls = false +disallow_untyped_defs = false ignore_missing_imports = true implicit_reexport = true no_implicit_optional = true From f3c650ec7c2cb45b161214b227eb4534bf88a3f4 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 28 Apr 2024 14:30:59 -0500 Subject: [PATCH 006/134] fix python-version string for mypy config --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index a2c7bf844..e1853a99a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -109,7 +109,7 @@ norecursedirs =["hooks", "*.egg", ".eggs", "dist", "build", "docs", ".tox", ".gi [tool.mypy] # https://mypy.readthedocs.io/en/latest/config_file.html#using-a-pyproject-toml-file -python_version = 3.10 +python_version = "3.10" pretty = true show_traceback = true color_output = true @@ -119,8 +119,8 @@ allow_redefinition = false check_untyped_defs = true disallow_any_generics = false disallow_incomplete_defs = true -disallow_untyped_calls = false -disallow_untyped_defs = false +disallow_untyped_calls = true +disallow_untyped_defs = true ignore_missing_imports = true implicit_reexport = true no_implicit_optional = true From 0af37942824580261830b621c76fe9bf0b7fb40a Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 28 Apr 2024 15:03:17 -0500 Subject: [PATCH 007/134] rename interface into protocol --- elastica/timestepper/__init__.py | 18 ++++-------------- .../{_stepper_interface.py => protocol.py} | 0 elastica/timestepper/symplectic_steppers.py | 10 ---------- tests/test_math/test_timestepper.py | 12 ------------ 4 files changed, 4 insertions(+), 36 deletions(-) rename elastica/timestepper/{_stepper_interface.py => protocol.py} (100%) diff --git a/elastica/timestepper/__init__.py b/elastica/timestepper/__init__.py index 3efa466eb..39ab84eb0 100644 --- a/elastica/timestepper/__init__.py +++ b/elastica/timestepper/__init__.py @@ -25,36 +25,26 @@ def extend_stepper_interface(Stepper, System): # by checking for the [] method is_this_system_a_collection = is_system_a_collection(System) - """ - # Stateful steppers are no more used so remove them - ConcreteStepper = ( - Stepper.stepper if _StatefulStepper in Stepper.__class__.mro() else Stepper - ) - """ - ConcreteStepper = Stepper - - if type(ConcreteStepper.Tag) == SymplecticStepperTag: + if type(Stepper.Tag) == SymplecticStepperTag: from elastica.timestepper.symplectic_steppers import ( _SystemInstanceStepper, _SystemCollectionStepper, SymplecticStepperMethods as StepperMethodCollector, ) - elif type(ConcreteStepper.Tag) == ExplicitStepperTag: + elif type(Stepper.Tag) == ExplicitStepperTag: from elastica.timestepper.explicit_steppers import ( _SystemInstanceStepper, _SystemCollectionStepper, ExplicitStepperMethods as StepperMethodCollector, ) - # elif SymplecticCosseratRodStepper in ConcreteStepper.__class__.mro(): - # return # hacky fix for now. remove HybridSteppers in a future version. else: raise NotImplementedError( "Only explicit and symplectic steppers are supported, given stepper is {}".format( - ConcreteStepper.__class__.__name__ + Stepper.__class__.__name__ ) ) - stepper_methods = StepperMethodCollector(ConcreteStepper) + stepper_methods = StepperMethodCollector(Stepper) do_step_method = ( _SystemCollectionStepper.do_step if is_this_system_a_collection diff --git a/elastica/timestepper/_stepper_interface.py b/elastica/timestepper/protocol.py similarity index 100% rename from elastica/timestepper/_stepper_interface.py rename to elastica/timestepper/protocol.py diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index 48dd42531..ae2a4ffe2 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -2,16 +2,6 @@ import numpy as np -# from elastica._elastica_numba._timestepper._symplectic_steppers import ( -# SymplecticStepperTag, -# PositionVerlet, -# PEFRL, -# ) - -# from elastica.timestepper._stepper_interface import ( -# _TimeStepper, -# _LinearExponentialIntegratorMixin, -# ) from elastica.rod.data_structures import ( overload_operator_kinematic_numba, overload_operator_dynamic_numba, diff --git a/tests/test_math/test_timestepper.py b/tests/test_math/test_timestepper.py index 878d92507..0511a99e7 100644 --- a/tests/test_math/test_timestepper.py +++ b/tests/test_math/test_timestepper.py @@ -14,7 +14,6 @@ ScalarExponentialDampedHarmonicOscillatorCollectiveSystem, ) from elastica.timestepper import integrate, extend_stepper_interface -from elastica.timestepper._stepper_interface import _TimeStepper from elastica.timestepper.explicit_steppers import ( RungeKutta4, @@ -159,17 +158,6 @@ def test_integrate_throws_an_assert_for_negative_total_steps(): SymplecticSteppers = [PositionVerlet, PEFRL] -class TestStepperInterface: - def test_no_base_access_error(self): - with pytest.raises(NotImplementedError) as excinfo: - _TimeStepper().do_step() - assert "not supposed to access" in str(excinfo.value) - - # @pytest.mark.parametrize("stepper", StatefulExplicitSteppers + SymplecticSteppers) - # def test_correct_orders(self, stepper): - # assert stepper().n_stages > 0, "Explicit stepper routine has no stages!" - - """ class TestExplicitSteppers: @pytest.mark.parametrize("stepper", StatefulExplicitSteppers) From afcbe6e0b3aa9c37db4e6eb32a3caefb5378471c Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 28 Apr 2024 16:15:42 -0500 Subject: [PATCH 008/134] change interface name to protocol --- elastica/timestepper/__init__.py | 3 --- elastica/timestepper/protocol.py | 15 ++++----------- elastica/timestepper/symplectic_steppers.py | 3 --- 3 files changed, 4 insertions(+), 17 deletions(-) diff --git a/elastica/timestepper/__init__.py b/elastica/timestepper/__init__.py index 39ab84eb0..6c4de76e7 100644 --- a/elastica/timestepper/__init__.py +++ b/elastica/timestepper/__init__.py @@ -18,7 +18,6 @@ # TODO: Both extend_stepper_interface and integrate should be in separate file. # __init__ is probably not an ideal place to have these scripts. def extend_stepper_interface(Stepper, System): - from elastica.utils import extend_instance from elastica.systems import is_system_a_collection # Check if system is a "collection" of smaller systems @@ -53,7 +52,6 @@ def extend_stepper_interface(Stepper, System): return do_step_method, stepper_methods.step_methods() -# TODO Improve interface of this function to take args and kwargs for ease of use def integrate( StatefulStepper, System, @@ -61,7 +59,6 @@ def integrate( n_steps: int = 1000, restart_time: float = 0.0, progress_bar: bool = True, - **kwargs, ): """ diff --git a/elastica/timestepper/protocol.py b/elastica/timestepper/protocol.py index 99d2a6ce8..e08baab95 100644 --- a/elastica/timestepper/protocol.py +++ b/elastica/timestepper/protocol.py @@ -1,17 +1,10 @@ __doc__ = "Time stepper interface" +from typing import Protocol -class _TimeStepper: - """Interface classes for all time-steppers""" - - def __init__(self): - pass - - def do_step(self, *args, **kwargs): - raise NotImplementedError( - "TimeStepper hierarchy is not supposed to access the do-step routine of the TimeStepper base class. " - ) - +class TimeStepperProtocol(Protocol): + """Protocol for all time-steppers""" + def do_step(self, *args, **kwargs): ... # class _StatefulStepper: # """ diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index ae2a4ffe2..3292433a4 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -202,9 +202,6 @@ class PositionVerlet: Tag = SymplecticStepperTag() - def __init__(self): - pass - def _first_prefactor(self, dt): return 0.5 * dt From 26f7c3bdeae8a08c08e4ae6b9291151a6a09ef33 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 28 Apr 2024 16:56:46 -0500 Subject: [PATCH 009/134] make separate tag module for easy typing --- elastica/timestepper/__init__.py | 7 +++--- elastica/timestepper/explicit_steppers.py | 16 ++++--------- elastica/timestepper/symplectic_steppers.py | 13 +++-------- elastica/timestepper/tag.py | 26 +++++++++++++++++++++ tests/test_math/test_timestepper.py | 7 +++--- 5 files changed, 40 insertions(+), 29 deletions(-) create mode 100644 elastica/timestepper/tag.py diff --git a/elastica/timestepper/__init__.py b/elastica/timestepper/__init__.py index 6c4de76e7..643707ce0 100644 --- a/elastica/timestepper/__init__.py +++ b/elastica/timestepper/__init__.py @@ -4,15 +4,14 @@ import numpy as np from tqdm import tqdm from elastica.timestepper.symplectic_steppers import ( - SymplecticStepperTag, PositionVerlet, PEFRL, ) from elastica.timestepper.explicit_steppers import ( - ExplicitStepperTag, RungeKutta4, EulerForward, ) +from elastica.timestepper.tag import SymplecticStepperTag, ExplicitStepperTag # TODO: Both extend_stepper_interface and integrate should be in separate file. @@ -24,13 +23,13 @@ def extend_stepper_interface(Stepper, System): # by checking for the [] method is_this_system_a_collection = is_system_a_collection(System) - if type(Stepper.Tag) == SymplecticStepperTag: + if isinstance(Stepper.Tag, SymplecticStepperTag): from elastica.timestepper.symplectic_steppers import ( _SystemInstanceStepper, _SystemCollectionStepper, SymplecticStepperMethods as StepperMethodCollector, ) - elif type(Stepper.Tag) == ExplicitStepperTag: + elif isinstance(Stepper.Tag, ExplicitStepperTag): from elastica.timestepper.explicit_steppers import ( _SystemInstanceStepper, _SystemCollectionStepper, diff --git a/elastica/timestepper/explicit_steppers.py b/elastica/timestepper/explicit_steppers.py index 2e06e9750..5e4e6ef48 100644 --- a/elastica/timestepper/explicit_steppers.py +++ b/elastica/timestepper/explicit_steppers.py @@ -2,6 +2,8 @@ import numpy as np from copy import copy +from .tag import tag, ExplicitStepperTag + """ Developer Note @@ -151,9 +153,8 @@ def n_stages(self): Classical EulerForward """ - +# @tag(ExplicitStepperTag) # class EulerForward: -# Tag = ExplicitStepperTag() # # def __init__(self): # pass @@ -188,20 +189,13 @@ def __init__(self): self.linear_operator = None """ - -class ExplicitStepperTag: - def __init__(self): - pass - - +@tag(ExplicitStepperTag) class RungeKutta4: """ Stateless runge-kutta4. coordinates operations only, memory needs to be externally managed and allocated. """ - Tag = ExplicitStepperTag() - def __init__(self): pass @@ -245,8 +239,8 @@ def _fourth_update(self, System, Memory, time: np.float64, dt: np.float64): return time +@tag(ExplicitStepperTag) class EulerForward: - Tag = ExplicitStepperTag() def __init__(self): super(EulerForward, self).__init__() diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index 3292433a4..b5f773bfa 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -6,6 +6,7 @@ overload_operator_kinematic_numba, overload_operator_dynamic_numba, ) +from .tag import tag, SymplecticStepperTag """ Developer Note @@ -189,19 +190,13 @@ def n_stages(self): return len(self._steps_and_prefactors) -class SymplecticStepperTag: - def __init__(self): - pass - - +@tag(SymplecticStepperTag) class PositionVerlet: """ Position Verlet symplectic time stepper class, which includes methods for second-order position Verlet. """ - Tag = SymplecticStepperTag() - def _first_prefactor(self, dt): return 0.5 * dt @@ -224,7 +219,7 @@ def _first_dynamic_step(self, System, time: np.float64, dt: np.float64): System.dynamic_rates(time, dt), ) - +@tag(SymplecticStepperTag) class PEFRL: """ Position Extended Forest-Ruth Like Algorithm of @@ -241,8 +236,6 @@ class PEFRL: lambda_dash_coeff = 0.5 * (1.0 - 2.0 * λ) xi_chi_dash_coeff = 1.0 - 2.0 * (ξ + χ) - Tag = SymplecticStepperTag() - def __init__(self): pass diff --git a/elastica/timestepper/tag.py b/elastica/timestepper/tag.py new file mode 100644 index 000000000..c8034d52d --- /dev/null +++ b/elastica/timestepper/tag.py @@ -0,0 +1,26 @@ +from typing import Type, List + +# TODO: Maybe move this for common utility +def tag(Tag) -> Type: + """ + Tag a class with arbitrary type-class + + example: + class ATag: ... + + @tag(ATag) + class A1: + ... + + assert isinstance(A1.tag, ATag) + """ + def wrapper(cls): + cls.Tag = Tag() + return cls + return wrapper + +class SymplecticStepperTag: + pass + +class ExplicitStepperTag: + pass diff --git a/tests/test_math/test_timestepper.py b/tests/test_math/test_timestepper.py index 0511a99e7..78c9cf605 100644 --- a/tests/test_math/test_timestepper.py +++ b/tests/test_math/test_timestepper.py @@ -17,7 +17,6 @@ from elastica.timestepper.explicit_steppers import ( RungeKutta4, - ExplicitStepperTag, EulerForward, ) @@ -28,8 +27,8 @@ from elastica.timestepper.symplectic_steppers import ( PositionVerlet, PEFRL, - SymplecticStepperTag, ) +from elastica.timestepper.tag import tag, SymplecticStepperTag, ExplicitStepperTag from elastica.utils import Tolerance @@ -38,8 +37,8 @@ class TestExtendStepperInterface: """TODO add documentation""" + @tag(SymplecticStepperTag) class MockSymplecticStepper: - Tag = SymplecticStepperTag() def _first_prefactor(self): pass @@ -50,8 +49,8 @@ def _first_kinematic_step(self): def _first_dynamic_step(self): pass + @tag(ExplicitStepperTag) class MockExplicitStepper: - Tag = ExplicitStepperTag() def _first_stage(self): pass From 016d69f186bde83f1aac3a1470ea93e0bddb27c5 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 28 Apr 2024 17:27:59 -0500 Subject: [PATCH 010/134] add system collection as native subtype --- elastica/typing.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/elastica/typing.py b/elastica/typing.py index d70a16433..9d12794c5 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -1,9 +1,11 @@ from elastica.rod import RodBase from elastica.rigidbody import RigidBodyBase from elastica.surface import SurfaceBase +from elastica.module import BaseSystemCollection from typing import Type, Union RodType = Type[RodBase] -SystemType = Union[RodType, Type[RigidBodyBase]] +SystemCollectionType = Type[BaseSystemCollection] +SystemType = Union[RodType, Type[RigidBodyBase], SystemCollectionType] AllowedContactType = Union[SystemType, Type[SurfaceBase]] From 237dc2f7a55ca095f848b44e820baeb824c1b212 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 28 Apr 2024 18:02:39 -0500 Subject: [PATCH 011/134] typehint is_system_a_collection --- elastica/systems/__init__.py | 2 +- pyproject.toml | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/elastica/systems/__init__.py b/elastica/systems/__init__.py index 33a39b836..16a17ff20 100644 --- a/elastica/systems/__init__.py +++ b/elastica/systems/__init__.py @@ -1,4 +1,4 @@ -def is_system_a_collection(system): +def is_system_a_collection(system: object) -> bool: # Check if system is a "collection" of smaller systems # by checking for the [] method """ diff --git a/pyproject.toml b/pyproject.toml index 68d482d6d..7cc5d1b08 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -132,9 +132,9 @@ strict_optional = true warn_no_return = true warn_redundant_casts = true warn_return_any = true -warn_unreachable = true +warn_unreachable = false warn_unused_configs = true -warn_unused_ignores = true +warn_unused_ignores = false [tool.coverage.report] # Regexes for lines to exclude from consideration From 63cf3426f2b5e586f36b030fe1af897950d1ae93 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 28 Apr 2024 18:05:29 -0500 Subject: [PATCH 012/134] wip: typing symplectic stepper --- elastica/timestepper/__init__.py | 63 +++++++++++---------- elastica/timestepper/explicit_steppers.py | 3 +- elastica/timestepper/protocol.py | 56 +++++++++++------- elastica/timestepper/symplectic_steppers.py | 7 ++- elastica/timestepper/tag.py | 5 ++ 5 files changed, 83 insertions(+), 51 deletions(-) diff --git a/elastica/timestepper/__init__.py b/elastica/timestepper/__init__.py index 643707ce0..f266d8eac 100644 --- a/elastica/timestepper/__init__.py +++ b/elastica/timestepper/__init__.py @@ -1,71 +1,76 @@ __doc__ = """Timestepping utilities to be used with Rod and RigidBody classes""" +from typing import Tuple, List, Callable, Type import numpy as np from tqdm import tqdm -from elastica.timestepper.symplectic_steppers import ( - PositionVerlet, - PEFRL, -) -from elastica.timestepper.explicit_steppers import ( - RungeKutta4, - EulerForward, -) -from elastica.timestepper.tag import SymplecticStepperTag, ExplicitStepperTag + +from elastica.typing import SystemType +from .symplectic_steppers import PositionVerlet, PEFRL +from .explicit_steppers import RungeKutta4, EulerForward +from .tag import SymplecticStepperTag, ExplicitStepperTag +from .protocol import StepperProtocol, StatefulStepperProtocol +from .protocol import MethodCollectorProtocol # TODO: Both extend_stepper_interface and integrate should be in separate file. # __init__ is probably not an ideal place to have these scripts. -def extend_stepper_interface(Stepper, System): +def extend_stepper_interface( + Stepper: StepperProtocol, System: SystemType +) -> Tuple[Callable, Tuple[Callable]]: from elastica.systems import is_system_a_collection - # Check if system is a "collection" of smaller systems - # by checking for the [] method - is_this_system_a_collection = is_system_a_collection(System) - + # StepperMethodCollector: Type[MethodCollectorProtocol] + # SystemStepper: Type[StepperProtocol] if isinstance(Stepper.Tag, SymplecticStepperTag): from elastica.timestepper.symplectic_steppers import ( _SystemInstanceStepper, _SystemCollectionStepper, - SymplecticStepperMethods as StepperMethodCollector, + SymplecticStepperMethods, ) - elif isinstance(Stepper.Tag, ExplicitStepperTag): + + StepperMethodCollector = SymplecticStepperMethods + elif isinstance(Stepper.Tag, ExplicitStepperTag): # type: ignore[no-redef] from elastica.timestepper.explicit_steppers import ( _SystemInstanceStepper, _SystemCollectionStepper, - ExplicitStepperMethods as StepperMethodCollector, + ExplicitStepperMethods, ) + + StepperMethodCollector = ExplicitStepperMethods else: raise NotImplementedError( "Only explicit and symplectic steppers are supported, given stepper is {}".format( Stepper.__class__.__name__ ) ) + # Check if system is a "collection" of smaller systems + # by checking for the [] method + if is_system_a_collection(System): + SystemStepper = _SystemCollectionStepper + else: + SystemStepper = _SystemInstanceStepper - stepper_methods = StepperMethodCollector(Stepper) - do_step_method = ( - _SystemCollectionStepper.do_step - if is_this_system_a_collection - else _SystemInstanceStepper.do_step - ) - return do_step_method, stepper_methods.step_methods() + stepper_methods: Tuple[Callable] = StepperMethodCollector(Stepper).step_methods() + do_step_method: Callable = SystemStepper.do_step + return do_step_method, stepper_methods def integrate( - StatefulStepper, - System, + StatefulStepper: StatefulStepperProtocol, + System: SystemType, final_time: float, n_steps: int = 1000, restart_time: float = 0.0, progress_bar: bool = True, -): +) -> float: """ Parameters ---------- - StatefulStepper : + StatefulStepper : StatefulStepperProtocol Stepper algorithm to use. - System : + System : SystemType The elastica-system to simulate. final_time : float Total simulation time. The timestep is determined by final_time / n_steps. diff --git a/elastica/timestepper/explicit_steppers.py b/elastica/timestepper/explicit_steppers.py index 5e4e6ef48..f127cb9b2 100644 --- a/elastica/timestepper/explicit_steppers.py +++ b/elastica/timestepper/explicit_steppers.py @@ -189,6 +189,7 @@ def __init__(self): self.linear_operator = None """ + @tag(ExplicitStepperTag) class RungeKutta4: """ @@ -239,7 +240,7 @@ def _fourth_update(self, System, Memory, time: np.float64, dt: np.float64): return time -@tag(ExplicitStepperTag) +@tag(ExplicitStepperTag) class EulerForward: def __init__(self): diff --git a/elastica/timestepper/protocol.py b/elastica/timestepper/protocol.py index e08baab95..01a2cad38 100644 --- a/elastica/timestepper/protocol.py +++ b/elastica/timestepper/protocol.py @@ -1,27 +1,45 @@ __doc__ = "Time stepper interface" -from typing import Protocol +from typing import Protocol, Tuple, Callable, Type -class TimeStepperProtocol(Protocol): +import numpy as np + + +class StepperProtocol(Protocol): """Protocol for all time-steppers""" - def do_step(self, *args, **kwargs): ... -# class _StatefulStepper: -# """ -# Stateful explicit, symplectic stepper wrapper. -# """ -# -# def __init__(self): -# pass -# -# # For stateful steppes, bind memory to self -# def do_step(self, System, time: np.float64, dt: np.float64): -# return self.stepper.do_step(System, self, time, dt) -# -# @property -# def n_stages(self): -# return self.stepper.n_stages -# + def do_step(self, *args, **kwargs) -> float: ... + + @property + def Tag(self) -> Type: ... + + +class StatefulStepperProtocol(StepperProtocol): + """ + Stateful explicit, symplectic stepper wrapper. + """ + + # For stateful steppes, bind memory to self + def do_step(self, System, time: np.floating, dt: np.floating) -> float: + """ + Perform one time step of the simulation. + Return the new time. + """ + ... + + @property + def n_stages(self) -> int: ... + + +class MethodCollectorProtocol(Protocol): + """ + Protocol for collecting stepper methods. + """ + + def __init__(self, timestepper_instance: StepperProtocol): ... + + def step_methods(self) -> Tuple[Callable]: ... + # class _LinearExponentialIntegratorMixin: # """ diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index 0810a7943..3a0081b68 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -1,5 +1,7 @@ __doc__ = """Symplectic time steppers and concepts for integrating the kinematic and dynamic equations of rod-like objects. """ +from typing import Tuple, Callable + import numpy as np from elastica.rod.data_structures import ( @@ -173,7 +175,7 @@ def mirror(in_list): def NoOp(*args): pass - self._steps_and_prefactors = tuple( + self._steps_and_prefactors: Tuple[Callable] = tuple( zip_longest( self._prefactors, self._kinematic_steps, @@ -182,7 +184,7 @@ def NoOp(*args): ) ) - def step_methods(self): + def step_methods(self) -> Tuple[Callable]: return self._steps_and_prefactors @property @@ -219,6 +221,7 @@ def _first_dynamic_step(self, System, time: np.float64, dt: np.float64): System.dynamic_rates(time, dt), ) + @tag(SymplecticStepperTag) class PEFRL: """ diff --git a/elastica/timestepper/tag.py b/elastica/timestepper/tag.py index c8034d52d..5d4c59ee8 100644 --- a/elastica/timestepper/tag.py +++ b/elastica/timestepper/tag.py @@ -1,5 +1,6 @@ from typing import Type, List + # TODO: Maybe move this for common utility def tag(Tag) -> Type: """ @@ -14,13 +15,17 @@ class A1: assert isinstance(A1.tag, ATag) """ + def wrapper(cls): cls.Tag = Tag() return cls + return wrapper + class SymplecticStepperTag: pass + class ExplicitStepperTag: pass From 538721c7d8672657d6e1ce55e5f2f086a4f3ca0c Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 28 Apr 2024 20:04:42 -0500 Subject: [PATCH 013/134] wip: timestepper typing --- elastica/timestepper/__init__.py | 13 +-- elastica/timestepper/symplectic_steppers.py | 112 +++++++++++--------- elastica/typing.py | 6 +- 3 files changed, 71 insertions(+), 60 deletions(-) diff --git a/elastica/timestepper/__init__.py b/elastica/timestepper/__init__.py index f266d8eac..88a1ccefb 100644 --- a/elastica/timestepper/__init__.py +++ b/elastica/timestepper/__init__.py @@ -6,6 +6,8 @@ from tqdm import tqdm from elastica.typing import SystemType +from elastica.systems import is_system_a_collection + from .symplectic_steppers import PositionVerlet, PEFRL from .explicit_steppers import RungeKutta4, EulerForward from .tag import SymplecticStepperTag, ExplicitStepperTag @@ -18,13 +20,11 @@ def extend_stepper_interface( Stepper: StepperProtocol, System: SystemType ) -> Tuple[Callable, Tuple[Callable]]: - from elastica.systems import is_system_a_collection # StepperMethodCollector: Type[MethodCollectorProtocol] # SystemStepper: Type[StepperProtocol] if isinstance(Stepper.Tag, SymplecticStepperTag): from elastica.timestepper.symplectic_steppers import ( - _SystemInstanceStepper, _SystemCollectionStepper, SymplecticStepperMethods, ) @@ -32,7 +32,6 @@ def extend_stepper_interface( StepperMethodCollector = SymplecticStepperMethods elif isinstance(Stepper.Tag, ExplicitStepperTag): # type: ignore[no-redef] from elastica.timestepper.explicit_steppers import ( - _SystemInstanceStepper, _SystemCollectionStepper, ExplicitStepperMethods, ) @@ -44,12 +43,10 @@ def extend_stepper_interface( Stepper.__class__.__name__ ) ) + # Check if system is a "collection" of smaller systems - # by checking for the [] method - if is_system_a_collection(System): - SystemStepper = _SystemCollectionStepper - else: - SystemStepper = _SystemInstanceStepper + assert is_system_a_collection(System) + SystemStepper = _SystemCollectionStepper stepper_methods: Tuple[Callable] = StepperMethodCollector(Stepper).step_methods() do_step_method: Callable = SystemStepper.do_step diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index 3a0081b68..ab8678c0c 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -1,6 +1,6 @@ __doc__ = """Symplectic time steppers and concepts for integrating the kinematic and dynamic equations of rod-like objects. """ -from typing import Tuple, Callable +from typing import Tuple, Callable, Any, List import numpy as np @@ -8,6 +8,8 @@ overload_operator_kinematic_numba, overload_operator_dynamic_numba, ) +from elastica.typing import SystemCollectionType as SystemType +from .protocol import StatefulStepperProtocol from .tag import tag, SymplecticStepperTag """ @@ -22,8 +24,12 @@ class _SystemInstanceStepper: @staticmethod def do_step( - TimeStepper, _steps_and_prefactors, System, time: np.float64, dt: np.float64 - ): + TimeStepper: StatefulStepperProtocol, + _steps_and_prefactors: Tuple[Tuple[Callable, Callable, Callable], ...], + System: SystemType, + time: np.floating, + dt: np.floating, + ) -> np.floating: for kin_prefactor, kin_step, dyn_step in _steps_and_prefactors[:-1]: kin_step(TimeStepper, System, time, dt) time += kin_prefactor(TimeStepper, dt) @@ -45,18 +51,18 @@ class _SystemCollectionStepper: @staticmethod def do_step( - TimeStepper, - _steps_and_prefactors, - SystemCollection, - time: np.float64, - dt: np.float64, - ): + TimeStepper: StatefulStepperProtocol, + _steps_and_prefactors: Tuple[Tuple[Callable, Callable, Callable], ...], + SystemCollection: SystemType, + time: np.floating, + dt: np.floating, + ) -> np.floating: """ Function for doing symplectic stepper over the user defined rods (system). Parameters ---------- - SystemCollection: rod object + SystemCollection: SystemType time: float dt: float @@ -108,20 +114,20 @@ def do_step( class SymplecticStepperMethods: - def __init__(self, timestepper_instance): + def __init__(self, timestepper_instance: StatefulStepperProtocol): take_methods_from = timestepper_instance # Let the total number of steps for the Symplectic method # be (2*n + 1) (for time-symmetry). What we do is collect # the first n + 1 entries down in _steps and _prefac below, and then # reverse and append it to itself. - self._steps = [ + self._steps: List[Callable] = [ v for (k, v) in take_methods_from.__class__.__dict__.items() if k.endswith("step") ] # Prefac here is necessary because the linear-exponential integrator # needs only the prefactor and not the dt. - self._prefactors = [ + self._prefactors: List[Callable] = [ v for (k, v) in take_methods_from.__class__.__dict__.items() if k.endswith("prefactor") @@ -135,7 +141,7 @@ def __init__(self, timestepper_instance): # if k.endswith("forces_torques") # ] - def mirror(in_list): + def mirror(in_list: List) -> None: """Mirrors an input list ignoring the last element If steps = [A, B, C] then this call makes it [A, B, C, B, A] @@ -144,9 +150,6 @@ def mirror(in_list): ---------- in_list : input list to be mirrored, modified in-place - Returns - ------- - """ # syntax is very ugly if len(in_list) > 1: @@ -161,8 +164,8 @@ def mirror(in_list): len(self._steps) == 2 * len(self._prefactors) - 1 ), "Size mismatch in the number of steps and prefactors provided for a Symplectic Stepper!" - self._kinematic_steps = self._steps[::2] - self._dynamic_steps = self._steps[1::2] + self._kinematic_steps: List[Callable] = self._steps[::2] + self._dynamic_steps: List[Callable] = self._steps[1::2] # Avoid this check for MockClasses if len(self._kinematic_steps) > 0: @@ -172,23 +175,25 @@ def mirror(in_list): from itertools import zip_longest - def NoOp(*args): + def NoOp(*args: Any) -> None: pass - self._steps_and_prefactors: Tuple[Callable] = tuple( - zip_longest( - self._prefactors, - self._kinematic_steps, - self._dynamic_steps, - fillvalue=NoOp, + self._steps_and_prefactors: Tuple[Tuple[Callable, Callable, Callable], ...] = ( + tuple( + zip_longest( + self._prefactors, + self._kinematic_steps, + self._dynamic_steps, + fillvalue=NoOp, + ) ) ) - def step_methods(self) -> Tuple[Callable]: + def step_methods(self) -> Tuple[Tuple[Callable, Callable, Callable], ...]: return self._steps_and_prefactors @property - def n_stages(self): + def n_stages(self) -> int: return len(self._steps_and_prefactors) @@ -199,12 +204,13 @@ class PositionVerlet: includes methods for second-order position Verlet. """ - def _first_prefactor(self, dt): + def _first_prefactor(self, dt: np.floating) -> np.floating: return 0.5 * dt - def _first_kinematic_step(self, System, time: np.float64, dt: np.float64): + def _first_kinematic_step( + self, System: SystemType, time: np.floating, dt: np.floating + ) -> None: prefac = self._first_prefactor(dt) - overload_operator_kinematic_numba( System.n_nodes, prefac, @@ -214,8 +220,9 @@ def _first_kinematic_step(self, System, time: np.float64, dt: np.float64): System.omega_collection, ) - def _first_dynamic_step(self, System, time: np.float64, dt: np.float64): - + def _first_dynamic_step( + self, System: SystemType, time: np.floating, dt: np.floating + ) -> None: overload_operator_dynamic_numba( System.dynamic_states.rate_collection, System.dynamic_rates(time, dt), @@ -231,21 +238,20 @@ class PEFRL: """ # xi and chi are confusing, but be careful! - ξ = np.float64(0.1786178958448091e0) # ξ - λ = -np.float64(0.2123418310626054e0) # λ - χ = -np.float64(0.6626458266981849e-1) # χ + ξ: np.float64 = np.float64(0.1786178958448091e0) # ξ + λ: np.float64 = -np.float64(0.2123418310626054e0) # λ + χ: np.float64 = -np.float64(0.6626458266981849e-1) # χ # Pre-calculate other coefficients - lambda_dash_coeff = 0.5 * (1.0 - 2.0 * λ) - xi_chi_dash_coeff = 1.0 - 2.0 * (ξ + χ) - - def __init__(self): - pass + lambda_dash_coeff: np.float64 = 0.5 * (1.0 - 2.0 * λ) + xi_chi_dash_coeff: np.float64 = 1.0 - 2.0 * (ξ + χ) - def _first_kinematic_prefactor(self, dt): + def _first_kinematic_prefactor(self, dt: np.floating) -> np.floating: return self.ξ * dt - def _first_kinematic_step(self, System, time: np.float64, dt: np.float64): + def _first_kinematic_step( + self, System: SystemType, time: np.floating, dt: np.floating + ) -> None: prefac = self._first_kinematic_prefactor(dt) overload_operator_kinematic_numba( System.n_nodes, @@ -257,7 +263,9 @@ def _first_kinematic_step(self, System, time: np.float64, dt: np.float64): ) # System.kinematic_states += prefac * System.kinematic_rates(time, prefac) - def _first_dynamic_step(self, System, time: np.float64, dt: np.float64): + def _first_dynamic_step( + self, System: SystemType, time: np.floating, dt: np.floating + ) -> None: prefac = self.lambda_dash_coeff * dt overload_operator_dynamic_numba( System.dynamic_states.rate_collection, @@ -265,10 +273,12 @@ def _first_dynamic_step(self, System, time: np.float64, dt: np.float64): ) # System.dynamic_states += prefac * System.dynamic_rates(time, prefac) - def _second_kinematic_prefactor(self, dt): + def _second_kinematic_prefactor(self, dt: np.floating) -> np.floating: return self.χ * dt - def _second_kinematic_step(self, System, time: np.float64, dt: np.float64): + def _second_kinematic_step( + self, System: SystemType, time: np.floating, dt: np.floating + ) -> None: prefac = self._second_kinematic_prefactor(dt) overload_operator_kinematic_numba( System.n_nodes, @@ -280,7 +290,9 @@ def _second_kinematic_step(self, System, time: np.float64, dt: np.float64): ) # System.kinematic_states += prefac * System.kinematic_rates(time, prefac) - def _second_dynamic_step(self, System, time: np.float64, dt: np.float64): + def _second_dynamic_step( + self, System: SystemType, time: np.floating, dt: np.floating + ) -> None: prefac = self.λ * dt overload_operator_dynamic_numba( System.dynamic_states.rate_collection, @@ -288,10 +300,12 @@ def _second_dynamic_step(self, System, time: np.float64, dt: np.float64): ) # System.dynamic_states += prefac * System.dynamic_rates(time, prefac) - def _third_kinematic_prefactor(self, dt): + def _third_kinematic_prefactor(self, dt: np.floating) -> np.floating: return self.xi_chi_dash_coeff * dt - def _third_kinematic_step(self, System, time: np.float64, dt: np.float64): + def _third_kinematic_step( + self, System: SystemType, time: np.floating, dt: np.floating + ) -> None: prefac = self._third_kinematic_prefactor(dt) # Need to fill in overload_operator_kinematic_numba( diff --git a/elastica/typing.py b/elastica/typing.py index 9d12794c5..e72c9307f 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -1,11 +1,11 @@ -from elastica.rod import RodBase +from elastica.rod import RodBase, RodProtocol from elastica.rigidbody import RigidBodyBase from elastica.surface import SurfaceBase from elastica.module import BaseSystemCollection from typing import Type, Union -RodType = Type[RodBase] +RodType = RodProtocol SystemCollectionType = Type[BaseSystemCollection] -SystemType = Union[RodType, Type[RigidBodyBase], SystemCollectionType] +SystemType = Union[RodProtocol, Type[RigidBodyBase], SystemCollectionType] AllowedContactType = Union[SystemType, Type[SurfaceBase]] From 6426bd1634b7217cf9d7ba8b2609ad59527d447e Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 28 Apr 2024 23:03:35 -0500 Subject: [PATCH 014/134] add generic system protocol --- elastica/systems/protocol.py | 29 +++++++++++++++++++++++++++++ elastica/typing.py | 10 +++++++--- 2 files changed, 36 insertions(+), 3 deletions(-) create mode 100644 elastica/systems/protocol.py diff --git a/elastica/systems/protocol.py b/elastica/systems/protocol.py new file mode 100644 index 000000000..397aefd35 --- /dev/null +++ b/elastica/systems/protocol.py @@ -0,0 +1,29 @@ +__doc__ = """Base class for elastica system""" + +from typing import Protocol + +from numpy.typing import NDArray + + +class SystemProtocol(Protocol): + """ + Protocol for all elastica system + """ + + @property + def position_collection(self) -> NDArray: ... + + @property + def omega_collection(self) -> NDArray: ... + + @property + def acceleration_collection(self) -> NDArray: ... + + @property + def alpha_collection(self) -> NDArray: ... + + @property + def external_forces(self) -> NDArray: ... + + @property + def external_torques(self) -> NDArray: ... diff --git a/elastica/typing.py b/elastica/typing.py index e72c9307f..fafc03bf7 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -1,11 +1,15 @@ -from elastica.rod import RodBase, RodProtocol +from elastica.rod import RodBase from elastica.rigidbody import RigidBodyBase from elastica.surface import SurfaceBase from elastica.module import BaseSystemCollection +from .system.protocol import SystemProtocol + from typing import Type, Union +from typing import Protocol + +SystemType = SystemProtocol -RodType = RodProtocol +RodType = Type[RodBase] SystemCollectionType = Type[BaseSystemCollection] -SystemType = Union[RodProtocol, Type[RigidBodyBase], SystemCollectionType] AllowedContactType = Union[SystemType, Type[SurfaceBase]] From df67eb4750d6e1053d5037473f7d1926f526fdb4 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 28 Apr 2024 23:44:28 -0500 Subject: [PATCH 015/134] wip explicit stepper typing --- elastica/timestepper/explicit_steppers.py | 230 ++++++++++++-------- elastica/timestepper/protocol.py | 15 +- elastica/timestepper/symplectic_steppers.py | 52 ++--- elastica/timestepper/tag.py | 6 +- elastica/typing.py | 10 +- 5 files changed, 187 insertions(+), 126 deletions(-) diff --git a/elastica/timestepper/explicit_steppers.py b/elastica/timestepper/explicit_steppers.py index f127cb9b2..2cf26f635 100644 --- a/elastica/timestepper/explicit_steppers.py +++ b/elastica/timestepper/explicit_steppers.py @@ -1,8 +1,13 @@ __doc__ = """Explicit timesteppers and concepts""" + +from typing import Any, Tuple + import numpy as np from copy import copy +from elastica.typing import SystemCollectionType, SystemType, SteppersOperatorsType from .tag import tag, ExplicitStepperTag +from .protocol import StatefulStepperProtocol, MemoryProtocol """ @@ -54,17 +59,21 @@ """ +# TODO: Modify this line and move to elastica/typing.py once system state is defined +StateType = Any + + class _SystemInstanceStepper: # # noinspection PyUnresolvedReferences @staticmethod def do_step( - TimeStepper, - _stages_and_updates, - System, - Memory, - time: np.float64, - dt: np.float64, - ): + TimeStepper: StatefulStepperProtocol, + _stages_and_updates: SteppersOperatorsType, + System: SystemType, + Memory: MemoryProtocol, + time: np.floating, + dt: np.floating, + ) -> np.floating: for stage, update in _stages_and_updates: stage(TimeStepper, System, Memory, time, dt) time = update(TimeStepper, System, Memory, time, dt) @@ -75,13 +84,13 @@ class _SystemCollectionStepper: # # noinspection PyUnresolvedReferences @staticmethod def do_step( - TimeStepper, - _stages_and_updates, - SystemCollection, - MemoryCollection, - time: np.float64, - dt: np.float64, - ): + TimeStepper: StatefulStepperProtocol, + _stages_and_updates: SteppersOperatorsType, + SystemCollection: SystemCollectionType, + MemoryCollection: Tuple[MemoryProtocol, ...], + time: np.floating, + dt: np.floating, + ) -> np.floating: for stage, update in _stages_and_updates: SystemCollection.synchronize(time) for system, memory in zip(SystemCollection[:-1], MemoryCollection[:-1]): @@ -100,7 +109,7 @@ class ExplicitStepperMethods: Can also be used as a mixin with optional cls argument below """ - def __init__(self, timestepper_instance): + def __init__(self, timestepper_instance: StatefulStepperProtocol): take_methods_from = timestepper_instance __stages = [ v @@ -123,71 +132,66 @@ def __init__(self, timestepper_instance): self._stages_and_updates = tuple(zip(__stages, __updates)) - def step_methods(self): + def step_methods(self) -> SteppersOperatorsType: return self._stages_and_updates @property - def n_stages(self): + def n_stages(self) -> int: return len(self._stages_and_updates) -# class StatefulRungeKutta4(_StatefulStepper): -# """ -# Stores all states of Rk within the time-stepper. Works as long as the states -# are all one big numpy array, made possible by carefully using views. -# -# Convenience wrapper around Stateless that provides memory -# """ -# -# def __init__(self): -# super(StatefulRungeKutta4, self).__init__() -# self.stepper = RungeKutta4() -# self.initial_state = None -# self.k_1 = None -# self.k_2 = None -# self.k_3 = None -# self.k_4 = None +class EulerForwardMemory: + def __init__(self, initial_state: StateType) -> None: + self.initial_state = initial_state -""" -Classical EulerForward -""" +@tag(ExplicitStepperTag) +class EulerForward: + """ + Classical Euler Forward stepper. Stateless, coordinates operations only. + """ -# @tag(ExplicitStepperTag) -# class EulerForward: -# -# def __init__(self): -# pass -# -# def _first_stage(self, System, Memory, time, dt): -# pass -# -# def _first_update(self, System, Memory, time, dt): -# System.state += dt * System(time, dt) -# return time + dt + def _first_stage( + self, + System: SystemType, + Memory: EulerForwardMemory, + time: np.floating, + dt: np.floating, + ) -> None: + pass + + def _first_update( + self, + System: SystemType, + Memory: EulerForwardMemory, + time: np.floating, + dt: np.floating, + ) -> np.floating: + System.state += dt * System(time, dt) + return time + dt -# class StatefulEulerForward(_StatefulStepper): -# def __init__(self): -# super(StatefulEulerForward, self).__init__() -# self.stepper = EulerForward() +class RungeKutta4Memory: + """ + Stores all states of Rk within the time-stepper. Works as long as the states + are all one big numpy array, made possible by carefully using views. + Convenience wrapper around Stateless that provides memory + """ -""" -class ExplicitLinearExponentialIntegrator( - _LinearExponentialIntegratorMixin, ExplicitStepper -): - def __init__(self): - _LinearExponentialIntegratorMixin.__init__(self) - ExplicitStepper.__init__(self, _LinearExponentialIntegratorMixin) - - -class StatefulLinearExponentialIntegrator(_StatefulStepper): - def __init__(self): - super(StatefulLinearExponentialIntegrator, self).__init__() - self.stepper = ExplicitLinearExponentialIntegrator() - self.linear_operator = None -""" + def __init__( + self, + initial_state: StateType, + k_1: StateType, + k_2: StateType, + k_3: StateType, + k_4: StateType, + ) -> None: + self.initial_state = initial_state + self.k_1 = k_1 + self.k_2 = k_2 + self.k_3 = k_3 + self.k_4 = k_4 @tag(ExplicitStepperTag) @@ -197,41 +201,86 @@ class RungeKutta4: to be externally managed and allocated. """ - def __init__(self): - pass - # These methods should be static, but because we need to enable automatic # discovery in ExplicitStepper, these are bound to the RungeKutta4 class # For automatic discovery, the order of declaring stages here is very important - def _first_stage(self, System, Memory, time: np.float64, dt: np.float64): + def _first_stage( + self, + System: SystemType, + Memory: RungeKutta4Memory, + time: np.floating, + dt: np.floating, + ) -> None: Memory.initial_state = copy(System.state) Memory.k_1 = dt * System(time, dt) # Don't update state yet - def _first_update(self, System, Memory, time: np.float64, dt: np.float64): + def _first_update( + self, + System: SystemType, + Memory: RungeKutta4Memory, + time: np.floating, + dt: np.floating, + ) -> np.floating: # prepare for next stage System.state = Memory.initial_state + 0.5 * Memory.k_1 return time + 0.5 * dt - def _second_stage(self, System, Memory, time: np.float64, dt: np.float64): + def _second_stage( + self, + System: SystemType, + Memory: RungeKutta4Memory, + time: np.floating, + dt: np.floating, + ) -> None: Memory.k_2 = dt * System(time, dt) # Don't update state yet - def _second_update(self, System, Memory, time: np.float64, dt: np.float64): + def _second_update( + self, + System: SystemType, + Memory: RungeKutta4Memory, + time: np.floating, + dt: np.floating, + ) -> np.floating: # prepare for next stage System.state = Memory.initial_state + 0.5 * Memory.k_2 return time - def _third_stage(self, System, Memory, time: np.float64, dt: np.float64): + def _third_stage( + self, + System: SystemType, + Memory: RungeKutta4Memory, + time: np.floating, + dt: np.floating, + ) -> None: Memory.k_3 = dt * System(time, dt) # Don't update state yet - def _third_update(self, System, Memory, time: np.float64, dt: np.float64): + def _third_update( + self, + System: SystemType, + Memory: RungeKutta4Memory, + time: np.floating, + dt: np.floating, + ) -> np.floating: # prepare for next stage System.state = Memory.initial_state + Memory.k_3 return time + 0.5 * dt - def _fourth_stage(self, System, Memory, time: np.float64, dt: np.float64): + def _fourth_stage( + self, + System: SystemType, + Memory: RungeKutta4Memory, + time: np.floating, + dt: np.floating, + ) -> None: Memory.k_4 = dt * System(time, dt) # Don't update state yet - def _fourth_update(self, System, Memory, time: np.float64, dt: np.float64): + def _fourth_update( + self, + System: SystemType, + Memory: RungeKutta4Memory, + time: np.floating, + dt: np.floating, + ) -> np.floating: # prepare for next stage System.state = ( Memory.initial_state @@ -240,15 +289,16 @@ def _fourth_update(self, System, Memory, time: np.float64, dt: np.float64): return time -@tag(ExplicitStepperTag) -class EulerForward: - - def __init__(self): - super(EulerForward, self).__init__() - - def _first_stage(self, System, Memory, time, dt): - pass - - def _first_update(self, System, Memory, time, dt): - System.state += dt * System(time, dt) - return time + dt +# class ExplicitLinearExponentialIntegrator( +# _LinearExponentialIntegratorMixin, ExplicitStepper +# ): +# def __init__(self): +# _LinearExponentialIntegratorMixin.__init__(self) +# ExplicitStepper.__init__(self, _LinearExponentialIntegratorMixin) +# +# +# class StatefulLinearExponentialIntegrator(_StatefulStepper): +# def __init__(self): +# super(StatefulLinearExponentialIntegrator, self).__init__() +# self.stepper = ExplicitLinearExponentialIntegrator() +# self.linear_operator = None diff --git a/elastica/timestepper/protocol.py b/elastica/timestepper/protocol.py index 01a2cad38..17e1e5918 100644 --- a/elastica/timestepper/protocol.py +++ b/elastica/timestepper/protocol.py @@ -1,26 +1,28 @@ __doc__ = "Time stepper interface" -from typing import Protocol, Tuple, Callable, Type +from typing import Protocol, Tuple, Callable, Type, Any import numpy as np +from elastica.typing import SystemType + class StepperProtocol(Protocol): """Protocol for all time-steppers""" - def do_step(self, *args, **kwargs) -> float: ... + def do_step(self, *args: Any, **kwargs: Any) -> float: ... @property def Tag(self) -> Type: ... -class StatefulStepperProtocol(StepperProtocol): +class StatefulStepperProtocol(StepperProtocol, Protocol): """ Stateful explicit, symplectic stepper wrapper. """ # For stateful steppes, bind memory to self - def do_step(self, System, time: np.floating, dt: np.floating) -> float: + def do_step(self, System: SystemType, time: np.floating, dt: np.floating) -> float: """ Perform one time step of the simulation. Return the new time. @@ -41,6 +43,11 @@ def __init__(self, timestepper_instance: StepperProtocol): ... def step_methods(self) -> Tuple[Callable]: ... +class MemoryProtocol(Protocol): + @property + def initial_state(self) -> bool: ... + + # class _LinearExponentialIntegratorMixin: # """ # Linear Exponential integrator mixin wrapper. diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index ab8678c0c..e6a6fab0b 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -8,7 +8,13 @@ overload_operator_kinematic_numba, overload_operator_dynamic_numba, ) -from elastica.typing import SystemCollectionType as SystemType +from elastica.typing import ( + SystemCollectionType, + SystemType, + StepOperatorType, + PrefactorOperatorType, + SteppersOperatorsType, +) from .protocol import StatefulStepperProtocol from .tag import tag, SymplecticStepperTag @@ -25,7 +31,7 @@ class _SystemInstanceStepper: @staticmethod def do_step( TimeStepper: StatefulStepperProtocol, - _steps_and_prefactors: Tuple[Tuple[Callable, Callable, Callable], ...], + _steps_and_prefactors: SteppersOperatorsType, System: SystemType, time: np.floating, dt: np.floating, @@ -37,7 +43,7 @@ def do_step( dyn_step(TimeStepper, System, time, dt) # Peel the last kinematic step and prefactor alone - last_kin_prefactor = _steps_and_prefactors[-1][0] + last_kin_prefactor: Callable[..., np.floating] = _steps_and_prefactors[-1][0] last_kin_step = _steps_and_prefactors[-1][1] last_kin_step(TimeStepper, System, time, dt) @@ -52,8 +58,8 @@ class _SystemCollectionStepper: @staticmethod def do_step( TimeStepper: StatefulStepperProtocol, - _steps_and_prefactors: Tuple[Tuple[Callable, Callable, Callable], ...], - SystemCollection: SystemType, + _steps_and_prefactors: SteppersOperatorsType, + SystemCollection: SystemCollectionType, time: np.floating, dt: np.floating, ) -> np.floating: @@ -62,7 +68,7 @@ def do_step( Parameters ---------- - SystemCollection: SystemType + SystemCollection: SystemCollectionType time: float dt: float @@ -120,28 +126,20 @@ def __init__(self, timestepper_instance: StatefulStepperProtocol): # be (2*n + 1) (for time-symmetry). What we do is collect # the first n + 1 entries down in _steps and _prefac below, and then # reverse and append it to itself. - self._steps: List[Callable] = [ + self._steps: List[StepOperatorType] = [ v for (k, v) in take_methods_from.__class__.__dict__.items() if k.endswith("step") ] # Prefac here is necessary because the linear-exponential integrator # needs only the prefactor and not the dt. - self._prefactors: List[Callable] = [ + self._prefactors: List[PrefactorOperatorType] = [ v for (k, v) in take_methods_from.__class__.__dict__.items() if k.endswith("prefactor") ] - # # We are getting function named as _update_internal_forces_torques from dictionary, - # # it turns a list. - # self._update_internal_forces_torques = [ - # v - # for (k, v) in take_methods_from.__class__.__dict__.items() - # if k.endswith("forces_torques") - # ] - - def mirror(in_list: List) -> None: + def mirror(in_list: List[Callable]) -> None: """Mirrors an input list ignoring the last element If steps = [A, B, C] then this call makes it [A, B, C, B, A] @@ -164,8 +162,8 @@ def mirror(in_list: List) -> None: len(self._steps) == 2 * len(self._prefactors) - 1 ), "Size mismatch in the number of steps and prefactors provided for a Symplectic Stepper!" - self._kinematic_steps: List[Callable] = self._steps[::2] - self._dynamic_steps: List[Callable] = self._steps[1::2] + self._kinematic_steps: List[StepOperatorType] = self._steps[::2] + self._dynamic_steps: List[StepOperatorType] = self._steps[1::2] # Avoid this check for MockClasses if len(self._kinematic_steps) > 0: @@ -178,18 +176,16 @@ def mirror(in_list: List) -> None: def NoOp(*args: Any) -> None: pass - self._steps_and_prefactors: Tuple[Tuple[Callable, Callable, Callable], ...] = ( - tuple( - zip_longest( - self._prefactors, - self._kinematic_steps, - self._dynamic_steps, - fillvalue=NoOp, - ) + self._steps_and_prefactors: SteppersOperatorsType = tuple( + zip_longest( + self._prefactors, + self._kinematic_steps, + self._dynamic_steps, + fillvalue=NoOp, ) ) - def step_methods(self) -> Tuple[Tuple[Callable, Callable, Callable], ...]: + def step_methods(self) -> SteppersOperatorsType: return self._steps_and_prefactors @property diff --git a/elastica/timestepper/tag.py b/elastica/timestepper/tag.py index 5d4c59ee8..158f0c66b 100644 --- a/elastica/timestepper/tag.py +++ b/elastica/timestepper/tag.py @@ -1,8 +1,8 @@ -from typing import Type, List +from typing import Type, List, Callable # TODO: Maybe move this for common utility -def tag(Tag) -> Type: +def tag(Tag: Type) -> Callable[[Type], Type]: """ Tag a class with arbitrary type-class @@ -16,7 +16,7 @@ class A1: assert isinstance(A1.tag, ATag) """ - def wrapper(cls): + def wrapper(cls: Type) -> Type: cls.Tag = Tag() return cls diff --git a/elastica/typing.py b/elastica/typing.py index fafc03bf7..7d5ed0de9 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -5,11 +5,19 @@ from .system.protocol import SystemProtocol -from typing import Type, Union +import numpy as np + +from typing import Type, Union, Callable, tuple from typing import Protocol SystemType = SystemProtocol +StepOperatorType = Callable[[SystemType, ...], None] +PrefactorOperatorType = Callable[[np.floating], np.floating] +SteppersOperatorsType = tuple[ + tuple[Union[PrefactorOperatorType, StepOperatorType], ...], ... +] + RodType = Type[RodBase] SystemCollectionType = Type[BaseSystemCollection] AllowedContactType = Union[SystemType, Type[SurfaceBase]] From 9b8266aa7a92e79e3c1347b4da880a030e20e69f Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Mon, 29 Apr 2024 00:15:33 -0500 Subject: [PATCH 016/134] test: fix rod/rigidbody mocking --- elastica/__init__.py | 1 - elastica/modules/base_system.py | 10 ++++++---- elastica/rod/rod_base.py | 6 ------ elastica/timestepper/__init__.py | 9 +++++++-- elastica/typing.py | 25 +++++++++++++++++-------- tests/test_contact_classes.py | 2 +- tests/test_contact_functions.py | 2 +- tests/test_contact_utils.py | 2 +- tests/test_joint.py | 3 ++- 9 files changed, 35 insertions(+), 25 deletions(-) diff --git a/elastica/__init__.py b/elastica/__init__.py index e96e00570..1cee3be48 100644 --- a/elastica/__init__.py +++ b/elastica/__init__.py @@ -79,7 +79,6 @@ ) from elastica._linalg import levi_civita_tensor from elastica.utils import isqrt -from elastica.typing import RodType, SystemType, AllowedContactType from elastica.timestepper import ( integrate, PositionVerlet, diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index 6dbf0bcc4..21a0ccbf3 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -7,6 +7,8 @@ """ from typing import Iterable, Callable, AnyStr +import numpy as np + from collections.abc import MutableSequence from elastica.rod import RodBase @@ -181,22 +183,22 @@ def finalize(self): self._feature_group_synchronize.pop(index) ) - def synchronize(self, time: float): + def synchronize(self, time: np.floating): # Collection call _feature_group_synchronize for feature in self._feature_group_synchronize: feature(time) - def constrain_values(self, time: float): + def constrain_values(self, time: np.floating): # Collection call _feature_group_constrain_values for feature in self._feature_group_constrain_values: feature(time) - def constrain_rates(self, time: float): + def constrain_rates(self, time: np.floating): # Collection call _feature_group_constrain_rates for feature in self._feature_group_constrain_rates: feature(time) - def apply_callbacks(self, time: float, current_step: int): + def apply_callbacks(self, time: np.floating, current_step: np.integer): # Collection call _feature_group_callback for feature in self._feature_group_callback: feature(time, current_step) diff --git a/elastica/rod/rod_base.py b/elastica/rod/rod_base.py index 621d13f07..15785a360 100644 --- a/elastica/rod/rod_base.py +++ b/elastica/rod/rod_base.py @@ -16,9 +16,3 @@ def __init__(self): RodBase does not take any arguments. """ pass - # self.position_collection = NotImplemented - # self.omega_collection = NotImplemented - # self.acceleration_collection = NotImplemented - # self.alpha_collection = NotImplemented - # self.external_forces = NotImplemented - # self.external_torques = NotImplemented diff --git a/elastica/timestepper/__init__.py b/elastica/timestepper/__init__.py index 88a1ccefb..9100036d7 100644 --- a/elastica/timestepper/__init__.py +++ b/elastica/timestepper/__init__.py @@ -10,6 +10,7 @@ from .symplectic_steppers import PositionVerlet, PEFRL from .explicit_steppers import RungeKutta4, EulerForward + from .tag import SymplecticStepperTag, ExplicitStepperTag from .protocol import StepperProtocol, StatefulStepperProtocol from .protocol import MethodCollectorProtocol @@ -25,6 +26,7 @@ def extend_stepper_interface( # SystemStepper: Type[StepperProtocol] if isinstance(Stepper.Tag, SymplecticStepperTag): from elastica.timestepper.symplectic_steppers import ( + _SystemInstanceStepper, _SystemCollectionStepper, SymplecticStepperMethods, ) @@ -32,6 +34,7 @@ def extend_stepper_interface( StepperMethodCollector = SymplecticStepperMethods elif isinstance(Stepper.Tag, ExplicitStepperTag): # type: ignore[no-redef] from elastica.timestepper.explicit_steppers import ( + _SystemInstanceStepper, _SystemCollectionStepper, ExplicitStepperMethods, ) @@ -45,8 +48,10 @@ def extend_stepper_interface( ) # Check if system is a "collection" of smaller systems - assert is_system_a_collection(System) - SystemStepper = _SystemCollectionStepper + if is_system_a_collection(System): + SystemStepper = _SystemCollectionStepper + else: + SystemStepper = _SystemInstanceStepper stepper_methods: Tuple[Callable] = StepperMethodCollector(Stepper).step_methods() do_step_method: Callable = SystemStepper.do_step diff --git a/elastica/typing.py b/elastica/typing.py index 7d5ed0de9..4154b1712 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -1,14 +1,23 @@ -from elastica.rod import RodBase -from elastica.rigidbody import RigidBodyBase -from elastica.surface import SurfaceBase -from elastica.module import BaseSystemCollection +from typing import TYPE_CHECKING +from typing import Type, Union, Callable +from typing import Protocol -from .system.protocol import SystemProtocol +from .systems.protocol import SystemProtocol import numpy as np -from typing import Type, Union, Callable, tuple -from typing import Protocol +if TYPE_CHECKING: + # Used for type hinting without circular imports + from .rod import RodBase + from .rigidbody import RigidBodyBase + from .surface import SurfaceBase + from .modules import BaseSystemCollection +else: + RodBase = None + RigidBodyBase = None + SurfaceBase = None + BaseSystemCollection = None + SystemType = SystemProtocol @@ -19,5 +28,5 @@ ] RodType = Type[RodBase] -SystemCollectionType = Type[BaseSystemCollection] +SystemCollectionType = BaseSystemCollection AllowedContactType = Union[SystemType, Type[SurfaceBase]] diff --git a/tests/test_contact_classes.py b/tests/test_contact_classes.py index 160b2fb71..94b41a457 100644 --- a/tests/test_contact_classes.py +++ b/tests/test_contact_classes.py @@ -12,7 +12,7 @@ RodPlaneContactWithAnisotropicFriction, CylinderPlaneContact, ) -from elastica.typing import RodBase +from elastica.rod import RodBase from elastica.rigidbody import Cylinder, Sphere from elastica.surface import Plane import pytest diff --git a/tests/test_contact_functions.py b/tests/test_contact_functions.py index f25d5830b..c3cae4b66 100644 --- a/tests/test_contact_functions.py +++ b/tests/test_contact_functions.py @@ -2,7 +2,7 @@ import numpy as np from numpy.testing import assert_allclose -from elastica.typing import RodBase +from elastica.rod import RodBase from elastica.rigidbody import Cylinder, Sphere from elastica._contact_functions import ( diff --git a/tests/test_contact_utils.py b/tests/test_contact_utils.py index e459d8154..4e5fe1bb8 100644 --- a/tests/test_contact_utils.py +++ b/tests/test_contact_utils.py @@ -4,7 +4,7 @@ import numpy as np from numpy.testing import assert_allclose from elastica.utils import Tolerance -from elastica.typing import RodBase +from elastica.rod import RodBase from elastica.rigidbody import Cylinder, Sphere from elastica.contact_utils import ( _dot_product, diff --git a/tests/test_joint.py b/tests/test_joint.py index d016226a2..373cd5f2d 100644 --- a/tests/test_joint.py +++ b/tests/test_joint.py @@ -480,7 +480,8 @@ def test_aabbs_not_intersecting_error_message(): assert error_info.value.args[0] == error_message -from elastica.typing import RodBase, RigidBodyBase +from elastica.rod import RodBase +from elastica.rigidbody import RigidBodyBase from elastica.joint import ( _prune_using_aabbs_rod_rigid_body, _prune_using_aabbs_rod_rod, From 577ff1ba200922f37fb2cfee8ee9f8ad59924895 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Mon, 29 Apr 2024 10:13:09 -0500 Subject: [PATCH 017/134] wip: resolve circular dependencies due to protocol --- elastica/modules/base_system.py | 2 +- elastica/rod/data_structures.py | 70 ++++++++++----------- elastica/systems/protocol.py | 52 ++++++++++++++- elastica/timestepper/__init__.py | 2 +- elastica/timestepper/explicit_steppers.py | 64 ++++++++++--------- elastica/timestepper/protocol.py | 4 +- elastica/timestepper/symplectic_steppers.py | 40 ++++++------ elastica/typing.py | 67 +++++++++++++++++--- 8 files changed, 200 insertions(+), 101 deletions(-) diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index 21a0ccbf3..34f42c7d4 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -198,7 +198,7 @@ def constrain_rates(self, time: np.floating): for feature in self._feature_group_constrain_rates: feature(time) - def apply_callbacks(self, time: np.floating, current_step: np.integer): + def apply_callbacks(self, time: np.floating, current_step: int): # Collection call _feature_group_callback for feature in self._feature_group_callback: feature(time, current_step) diff --git a/elastica/rod/data_structures.py b/elastica/rod/data_structures.py index 2a9466eb6..3c8b40328 100644 --- a/elastica/rod/data_structures.py +++ b/elastica/rod/data_structures.py @@ -8,40 +8,38 @@ # FIXME : Explicit Stepper doesn't work as States lose the # views they initially had when working with a timestepper. -""" -class _RodExplicitStepperMixin: - def __init__(self): - ( - self.state, - self.__deriv_state, - self.position_collection, - self.director_collection, - self.velocity_collection, - self.omega_collection, - self.acceleration_collection, - self.alpha_collection, # angular acceleration - ) = _bootstrap_from_data( - "explicit", self.n_elems, self._vector_states, self._matrix_states - ) - - # def __setattr__(self, name, value): - # np.copy(self.__dict__[name], value) - - def __call__(self, time, *args, **kwargs): - self.update_accelerations(time) # Internal, external - - # print("KRC", self.state.kinematic_rate_collection) - # print("DEr", self.__deriv_state.rate_collection) - if np.shares_memory( - self.state.kinematic_rate_collection, - self.velocity_collection - # self.__deriv_state.rate_collection - ): - print("Shares memory") - else: - print("Explicit states does not share memory") - return self.__deriv_state -""" +# class _RodExplicitStepperMixin: +# def __init__(self): +# ( +# self.state, +# self.__deriv_state, +# self.position_collection, +# self.director_collection, +# self.velocity_collection, +# self.omega_collection, +# self.acceleration_collection, +# self.alpha_collection, # angular acceleration +# ) = _bootstrap_from_data( +# "explicit", self.n_elems, self._vector_states, self._matrix_states +# ) +# +# # def __setattr__(self, name, value): +# # np.copy(self.__dict__[name], value) +# +# def __call__(self, time, *args, **kwargs): +# self.update_accelerations(time) # Internal, external +# +# # print("KRC", self.state.kinematic_rate_collection) +# # print("DEr", self.__deriv_state.rate_collection) +# if np.shares_memory( +# self.state.kinematic_rate_collection, +# self.velocity_collection +# # self.__deriv_state.rate_collection +# ): +# print("Shares memory") +# else: +# print("Explicit states does not share memory") +# return self.__deriv_state class _RodSymplecticStepperMixin: @@ -472,7 +470,7 @@ def __init__( self.velocity_collection = velocity_collection self.omega_collection = omega_collection - def kinematic_rates(self, time, *args, **kwargs): + def kinematic_rates(self, time, prefac): """Yields kinematic rates to interact with _KinematicState Returns @@ -488,7 +486,7 @@ def kinematic_rates(self, time, *args, **kwargs): # Comes from kin_state -> (x,Q) += dt * (v,w) <- First part of dyn_state return self.velocity_collection, self.omega_collection - def dynamic_rates(self, time, prefac, *args, **kwargs): + def dynamic_rates(self, time, prefac): """Yields dynamic rates to add to with _DynamicState Returns ------- diff --git a/elastica/systems/protocol.py b/elastica/systems/protocol.py index 397aefd35..a9462536a 100644 --- a/elastica/systems/protocol.py +++ b/elastica/systems/protocol.py @@ -1,7 +1,11 @@ __doc__ = """Base class for elastica system""" -from typing import Protocol +from typing import Protocol, Any +from elastica.typing import StateType +from elastica.rod.data_structures import _KinematicState, _DynamicState + +import numpy as np from numpy.typing import NDArray @@ -10,15 +14,21 @@ class SystemProtocol(Protocol): Protocol for all elastica system """ + @property + def n_nodes(self) -> int: ... + @property def position_collection(self) -> NDArray: ... @property - def omega_collection(self) -> NDArray: ... + def velocity_collection(self) -> NDArray: ... @property def acceleration_collection(self) -> NDArray: ... + @property + def omega_collection(self) -> NDArray: ... + @property def alpha_collection(self) -> NDArray: ... @@ -27,3 +37,41 @@ def external_forces(self) -> NDArray: ... @property def external_torques(self) -> NDArray: ... + + +class SymplecticSystemProtocol(SystemProtocol, Protocol): + """ + Protocol for system with symplectic state variables + """ + + @property + def kinematic_states(self) -> _KinematicState: ... + + @property + def dynamic_states(self) -> _DynamicState: ... + + @property + def rate_collection(self) -> NDArray: ... + + @property + def dvdt_dwdt_collection(self) -> NDArray: ... + + def kinematic_rates( + self, time: np.floating, prefac: np.floating + ) -> tuple[NDArray, NDArray]: ... + + def dynamic_rates( + self, time: np.floating, prefac: np.floating + ) -> tuple[NDArray]: ... + + def update_internal_forces_and_torques(self, time: np.floating) -> None: ... + + +class ExplicitSystemProtocol(SystemProtocol, Protocol): + # TODO: Temporarily made to handle explicit stepper. + # Need to be refactored as the explicit stepper is further developed. + def __call__(self, time: np.floating, dt: np.floating) -> np.floating: ... + @property + def state(self) -> StateType: ... + @state.setter + def state(self, state: StateType) -> None: ... diff --git a/elastica/timestepper/__init__.py b/elastica/timestepper/__init__.py index 9100036d7..66f8190fe 100644 --- a/elastica/timestepper/__init__.py +++ b/elastica/timestepper/__init__.py @@ -1,11 +1,11 @@ __doc__ = """Timestepping utilities to be used with Rod and RigidBody classes""" from typing import Tuple, List, Callable, Type +from elastica.typing import SystemType import numpy as np from tqdm import tqdm -from elastica.typing import SystemType from elastica.systems import is_system_a_collection from .symplectic_steppers import PositionVerlet, PEFRL diff --git a/elastica/timestepper/explicit_steppers.py b/elastica/timestepper/explicit_steppers.py index 2cf26f635..ca21f8dea 100644 --- a/elastica/timestepper/explicit_steppers.py +++ b/elastica/timestepper/explicit_steppers.py @@ -5,7 +5,13 @@ import numpy as np from copy import copy -from elastica.typing import SystemCollectionType, SystemType, SteppersOperatorsType +from elastica.typing import ( + SystemCollectionType, + OperatorType, + ExplicitOperatorsType, + StateType, +) +from elastica.systems.protocol import ExplicitSystemProtocol from .tag import tag, ExplicitStepperTag from .protocol import StatefulStepperProtocol, MemoryProtocol @@ -59,17 +65,13 @@ """ -# TODO: Modify this line and move to elastica/typing.py once system state is defined -StateType = Any - - class _SystemInstanceStepper: # # noinspection PyUnresolvedReferences @staticmethod def do_step( TimeStepper: StatefulStepperProtocol, - _stages_and_updates: SteppersOperatorsType, - System: SystemType, + _stages_and_updates: ExplicitOperatorsType, + System: ExplicitSystemProtocol, Memory: MemoryProtocol, time: np.floating, dt: np.floating, @@ -85,7 +87,7 @@ class _SystemCollectionStepper: @staticmethod def do_step( TimeStepper: StatefulStepperProtocol, - _stages_and_updates: SteppersOperatorsType, + _stages_and_updates: ExplicitOperatorsType, SystemCollection: SystemCollectionType, MemoryCollection: Tuple[MemoryProtocol, ...], time: np.floating, @@ -111,20 +113,20 @@ class ExplicitStepperMethods: def __init__(self, timestepper_instance: StatefulStepperProtocol): take_methods_from = timestepper_instance - __stages = [ + __stages: list[OperatorType] = [ v for (k, v) in take_methods_from.__class__.__dict__.items() if k.endswith("stage") ] - __updates = [ + __updates: list[OperatorType] = [ v for (k, v) in take_methods_from.__class__.__dict__.items() if k.endswith("update") ] # Tuples are almost immutable - _n_stages = len(__stages) - _n_updates = len(__updates) + _n_stages: int = len(__stages) + _n_updates: int = len(__updates) assert ( _n_stages == _n_updates @@ -132,7 +134,7 @@ def __init__(self, timestepper_instance: StatefulStepperProtocol): self._stages_and_updates = tuple(zip(__stages, __updates)) - def step_methods(self) -> SteppersOperatorsType: + def step_methods(self) -> ExplicitOperatorsType: return self._stages_and_updates @property @@ -153,7 +155,7 @@ class EulerForward: def _first_stage( self, - System: SystemType, + System: ExplicitSystemProtocol, Memory: EulerForwardMemory, time: np.floating, dt: np.floating, @@ -162,7 +164,7 @@ def _first_stage( def _first_update( self, - System: SystemType, + System: ExplicitSystemProtocol, Memory: EulerForwardMemory, time: np.floating, dt: np.floating, @@ -206,77 +208,77 @@ class RungeKutta4: # For automatic discovery, the order of declaring stages here is very important def _first_stage( self, - System: SystemType, + System: ExplicitSystemProtocol, Memory: RungeKutta4Memory, time: np.floating, dt: np.floating, ) -> None: Memory.initial_state = copy(System.state) - Memory.k_1 = dt * System(time, dt) # Don't update state yet + Memory.k_1 = dt * System(time, dt) # type: ignore[operator, assignment] def _first_update( self, - System: SystemType, + System: ExplicitSystemProtocol, Memory: RungeKutta4Memory, time: np.floating, dt: np.floating, ) -> np.floating: # prepare for next stage - System.state = Memory.initial_state + 0.5 * Memory.k_1 + System.state = Memory.initial_state + 0.5 * Memory.k_1 # type: ignore[operator] return time + 0.5 * dt def _second_stage( self, - System: SystemType, + System: ExplicitSystemProtocol, Memory: RungeKutta4Memory, time: np.floating, dt: np.floating, ) -> None: - Memory.k_2 = dt * System(time, dt) # Don't update state yet + Memory.k_2 = dt * System(time, dt) # type: ignore[operator, assignment] def _second_update( self, - System: SystemType, + System: ExplicitSystemProtocol, Memory: RungeKutta4Memory, time: np.floating, dt: np.floating, ) -> np.floating: # prepare for next stage - System.state = Memory.initial_state + 0.5 * Memory.k_2 + System.state = Memory.initial_state + 0.5 * Memory.k_2 # type: ignore[operator] return time def _third_stage( self, - System: SystemType, + System: ExplicitSystemProtocol, Memory: RungeKutta4Memory, time: np.floating, dt: np.floating, ) -> None: - Memory.k_3 = dt * System(time, dt) # Don't update state yet + Memory.k_3 = dt * System(time, dt) # type: ignore[operator, assignment] def _third_update( self, - System: SystemType, + System: ExplicitSystemProtocol, Memory: RungeKutta4Memory, time: np.floating, dt: np.floating, ) -> np.floating: # prepare for next stage - System.state = Memory.initial_state + Memory.k_3 + System.state = Memory.initial_state + Memory.k_3 # type: ignore[operator] return time + 0.5 * dt def _fourth_stage( self, - System: SystemType, + System: ExplicitSystemProtocol, Memory: RungeKutta4Memory, time: np.floating, dt: np.floating, ) -> None: - Memory.k_4 = dt * System(time, dt) # Don't update state yet + Memory.k_4 = dt * System(time, dt) # type: ignore[operator, assignment] def _fourth_update( self, - System: SystemType, + System: ExplicitSystemProtocol, Memory: RungeKutta4Memory, time: np.floating, dt: np.floating, @@ -284,7 +286,7 @@ def _fourth_update( # prepare for next stage System.state = ( Memory.initial_state - + (Memory.k_1 + 2.0 * Memory.k_2 + 2.0 * Memory.k_3 + Memory.k_4) / 6.0 + + (Memory.k_1 + 2.0 * Memory.k_2 + 2.0 * Memory.k_3 + Memory.k_4) / 6.0 # type: ignore[operator] ) return time diff --git a/elastica/timestepper/protocol.py b/elastica/timestepper/protocol.py index 17e1e5918..63a9a52d6 100644 --- a/elastica/timestepper/protocol.py +++ b/elastica/timestepper/protocol.py @@ -2,10 +2,10 @@ from typing import Protocol, Tuple, Callable, Type, Any -import numpy as np - from elastica.typing import SystemType +import numpy as np + class StepperProtocol(Protocol): """Protocol for all time-steppers""" diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index e6a6fab0b..783b34cc6 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -2,19 +2,21 @@ from typing import Tuple, Callable, Any, List +from elastica.typing import ( + SystemCollectionType, + # StepOperatorType, + # PrefactorOperatorType, + OperatorType, + SteppersOperatorsType, +) + import numpy as np from elastica.rod.data_structures import ( overload_operator_kinematic_numba, overload_operator_dynamic_numba, ) -from elastica.typing import ( - SystemCollectionType, - SystemType, - StepOperatorType, - PrefactorOperatorType, - SteppersOperatorsType, -) +from elastica.systems.protocol import SymplecticSystemProtocol from .protocol import StatefulStepperProtocol from .tag import tag, SymplecticStepperTag @@ -32,7 +34,7 @@ class _SystemInstanceStepper: def do_step( TimeStepper: StatefulStepperProtocol, _steps_and_prefactors: SteppersOperatorsType, - System: SystemType, + System: SymplecticSystemProtocol, time: np.floating, dt: np.floating, ) -> np.floating: @@ -126,14 +128,14 @@ def __init__(self, timestepper_instance: StatefulStepperProtocol): # be (2*n + 1) (for time-symmetry). What we do is collect # the first n + 1 entries down in _steps and _prefac below, and then # reverse and append it to itself. - self._steps: List[StepOperatorType] = [ + self._steps: List[OperatorType] = [ v for (k, v) in take_methods_from.__class__.__dict__.items() if k.endswith("step") ] # Prefac here is necessary because the linear-exponential integrator # needs only the prefactor and not the dt. - self._prefactors: List[PrefactorOperatorType] = [ + self._prefactors: List[OperatorType] = [ v for (k, v) in take_methods_from.__class__.__dict__.items() if k.endswith("prefactor") @@ -162,8 +164,8 @@ def mirror(in_list: List[Callable]) -> None: len(self._steps) == 2 * len(self._prefactors) - 1 ), "Size mismatch in the number of steps and prefactors provided for a Symplectic Stepper!" - self._kinematic_steps: List[StepOperatorType] = self._steps[::2] - self._dynamic_steps: List[StepOperatorType] = self._steps[1::2] + self._kinematic_steps: List[OperatorType] = self._steps[::2] + self._dynamic_steps: List[OperatorType] = self._steps[1::2] # Avoid this check for MockClasses if len(self._kinematic_steps) > 0: @@ -204,7 +206,7 @@ def _first_prefactor(self, dt: np.floating) -> np.floating: return 0.5 * dt def _first_kinematic_step( - self, System: SystemType, time: np.floating, dt: np.floating + self, System: SymplecticSystemProtocol, time: np.floating, dt: np.floating ) -> None: prefac = self._first_prefactor(dt) overload_operator_kinematic_numba( @@ -217,7 +219,7 @@ def _first_kinematic_step( ) def _first_dynamic_step( - self, System: SystemType, time: np.floating, dt: np.floating + self, System: SymplecticSystemProtocol, time: np.floating, dt: np.floating ) -> None: overload_operator_dynamic_numba( System.dynamic_states.rate_collection, @@ -246,7 +248,7 @@ def _first_kinematic_prefactor(self, dt: np.floating) -> np.floating: return self.ξ * dt def _first_kinematic_step( - self, System: SystemType, time: np.floating, dt: np.floating + self, System: SymplecticSystemProtocol, time: np.floating, dt: np.floating ) -> None: prefac = self._first_kinematic_prefactor(dt) overload_operator_kinematic_numba( @@ -260,7 +262,7 @@ def _first_kinematic_step( # System.kinematic_states += prefac * System.kinematic_rates(time, prefac) def _first_dynamic_step( - self, System: SystemType, time: np.floating, dt: np.floating + self, System: SymplecticSystemProtocol, time: np.floating, dt: np.floating ) -> None: prefac = self.lambda_dash_coeff * dt overload_operator_dynamic_numba( @@ -273,7 +275,7 @@ def _second_kinematic_prefactor(self, dt: np.floating) -> np.floating: return self.χ * dt def _second_kinematic_step( - self, System: SystemType, time: np.floating, dt: np.floating + self, System: SymplecticSystemProtocol, time: np.floating, dt: np.floating ) -> None: prefac = self._second_kinematic_prefactor(dt) overload_operator_kinematic_numba( @@ -287,7 +289,7 @@ def _second_kinematic_step( # System.kinematic_states += prefac * System.kinematic_rates(time, prefac) def _second_dynamic_step( - self, System: SystemType, time: np.floating, dt: np.floating + self, System: SymplecticSystemProtocol, time: np.floating, dt: np.floating ) -> None: prefac = self.λ * dt overload_operator_dynamic_numba( @@ -300,7 +302,7 @@ def _third_kinematic_prefactor(self, dt: np.floating) -> np.floating: return self.xi_chi_dash_coeff * dt def _third_kinematic_step( - self, System: SystemType, time: np.floating, dt: np.floating + self, System: SymplecticSystemProtocol, time: np.floating, dt: np.floating ) -> None: prefac = self._third_kinematic_prefactor(dt) # Need to fill in diff --git a/elastica/typing.py b/elastica/typing.py index 4154b1712..b23d5b602 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -1,31 +1,80 @@ +__doc__ = """ +This module contains aliases of type-hints for elastica. + +""" + from typing import TYPE_CHECKING -from typing import Type, Union, Callable +from typing import Type, Union, Callable, Any, VarArg from typing import Protocol -from .systems.protocol import SystemProtocol - import numpy as np + if TYPE_CHECKING: # Used for type hinting without circular imports + # NEVER BACK-IMPORT ANY ELASTICA MODULES HERE from .rod import RodBase from .rigidbody import RigidBodyBase from .surface import SurfaceBase from .modules import BaseSystemCollection + + from .rod.data_structures import _State as State + from .systems.protocol import SymplecticSystemProtocol, ExplicitSystemProtocol + from .timestepper.protocol import StatefulStepperProtocol, MemoryProtocol else: RodBase = None RigidBodyBase = None SurfaceBase = None BaseSystemCollection = None + State = "State" + SymplecticSystemProtocol = None + ExplicitSystemProtocol = None + StatefulStepperProtocol = None + MemoryProtocol = None + + +SystemType = Union[SymplecticSystemProtocol, ExplicitSystemProtocol] -SystemType = SystemProtocol +# TODO: Modify this line and move to elastica/typing.py once system state is defined +# Mostly used in explicit stepper: for symplectic, use kinetic and dynamic state +StateType = State -StepOperatorType = Callable[[SystemType, ...], None] -PrefactorOperatorType = Callable[[np.floating], np.floating] -SteppersOperatorsType = tuple[ - tuple[Union[PrefactorOperatorType, StepOperatorType], ...], ... -] +# NoOpt stepper +# Symplectic stepper +# StepOperatorType = Callable[ +# [StatefulStepperProtocol, SymplecticSystemProtocol, np.floating, np.floating], None +# ] +# PrefactorOperatorType = Callable[ +# [StatefulStepperProtocol, np.floating], np.floating +# ] +OperatorType = Callable[ + [VarArg(Any)], np.floating | None +] # TODO: Maybe can be more specific. Up for discussion. +SteppersOperatorsType = tuple[tuple[OperatorType, ...], ...] +# tuple[Union[PrefactorOperatorType, StepOperatorType, NoOpType, np.floating], ...], ... +# Explicit stepper +# ExplicitStageOperatorType = Callable[ +# [ +# StatefulStepperProtocol, +# ExplicitSystemProtocol, +# MemoryProtocol, +# np.floating, +# np.floating, +# ], +# None, +# ] +# ExplicitUpdateOperatorType = Callable[ +# [ +# StatefulStepperProtocol, +# ExplicitSystemProtocol, +# MemoryProtocol, +# np.floating, +# np.floating, +# ], +# np.floating, +# ] +ExplicitOperatorsType = tuple[tuple[OperatorType, ...], ...] RodType = Type[RodBase] SystemCollectionType = BaseSystemCollection From f0f5c5ad47b3997e1fae1f39973ade48932145a2 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Mon, 29 Apr 2024 10:41:47 -0500 Subject: [PATCH 018/134] remove protocol from coverage --- elastica/boundary_conditions.py | 4 ++-- elastica/modules/connections.py | 4 ++-- elastica/typing.py | 22 +++++++++++----------- pyproject.toml | 5 ++++- 4 files changed, 19 insertions(+), 16 deletions(-) diff --git a/elastica/boundary_conditions.py b/elastica/boundary_conditions.py index adca8b6e3..0cf666c8c 100644 --- a/elastica/boundary_conditions.py +++ b/elastica/boundary_conditions.py @@ -316,12 +316,12 @@ def __init__( rotational_constraint_selector = np.array([True, True, True]) # properly validate the user-provided constraint selectors assert ( - type(translational_constraint_selector) == np.ndarray + isinstance(translational_constraint_selector, np.ndarray) and translational_constraint_selector.dtype == bool and translational_constraint_selector.shape == (3,) ), "Translational constraint selector must be a 1D boolean array of length 3." assert ( - type(rotational_constraint_selector) == np.ndarray + isinstance(rotational_constraint_selector, np.ndarray) and rotational_constraint_selector.dtype == bool and rotational_constraint_selector.shape == (3,) ), "Rotational constraint selector must be a 1D boolean array of length 3." diff --git a/elastica/modules/connections.py b/elastica/modules/connections.py index 8c6cd6267..29e17ca03 100644 --- a/elastica/modules/connections.py +++ b/elastica/modules/connections.py @@ -152,8 +152,8 @@ def __init__( def set_index(self, first_idx, second_idx): # TODO assert range # First check if the types of first rod idx and second rod idx variable are same. - assert type(first_idx) == type( - second_idx + assert isinstance( + first_idx, type(second_idx) ), "Type of first_connect_idx :{}".format( type(first_idx) ) + " is different than second_connect_idx :{}".format( diff --git a/elastica/typing.py b/elastica/typing.py index b23d5b602..6334bfd4c 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -4,8 +4,8 @@ """ from typing import TYPE_CHECKING -from typing import Type, Union, Callable, Any, VarArg -from typing import Protocol +from typing import Type, Union, Callable, Any +from typing import Protocol, TypeAlias import numpy as np @@ -34,11 +34,11 @@ MemoryProtocol = None -SystemType = Union[SymplecticSystemProtocol, ExplicitSystemProtocol] +SystemType: TypeAlias = Union[SymplecticSystemProtocol, ExplicitSystemProtocol] # TODO: Modify this line and move to elastica/typing.py once system state is defined # Mostly used in explicit stepper: for symplectic, use kinetic and dynamic state -StateType = State +StateType: TypeAlias = State # NoOpt stepper # Symplectic stepper @@ -48,10 +48,10 @@ # PrefactorOperatorType = Callable[ # [StatefulStepperProtocol, np.floating], np.floating # ] -OperatorType = Callable[ - [VarArg(Any)], np.floating | None +OperatorType: TypeAlias = Callable[ + Any, Any ] # TODO: Maybe can be more specific. Up for discussion. -SteppersOperatorsType = tuple[tuple[OperatorType, ...], ...] +SteppersOperatorsType: TypeAlias = tuple[tuple[OperatorType, ...], ...] # tuple[Union[PrefactorOperatorType, StepOperatorType, NoOpType, np.floating], ...], ... # Explicit stepper # ExplicitStageOperatorType = Callable[ @@ -74,8 +74,8 @@ # ], # np.floating, # ] -ExplicitOperatorsType = tuple[tuple[OperatorType, ...], ...] +ExplicitOperatorsType: TypeAlias = tuple[tuple[OperatorType, ...], ...] -RodType = Type[RodBase] -SystemCollectionType = BaseSystemCollection -AllowedContactType = Union[SystemType, Type[SurfaceBase]] +RodType: TypeAlias = Type[RodBase] +SystemCollectionType: TypeAlias = BaseSystemCollection +AllowedContactType: TypeAlias = Union[SystemType, Type[SurfaceBase]] diff --git a/pyproject.toml b/pyproject.toml index 7cc5d1b08..d8988a3c5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -148,13 +148,15 @@ exclude_lines = [ "if __name__ == __main__:", "pass", "def __repr__", +#"if self.debug:", +#"if settings.DEBUG", "from", "import", "if TYPE_CHECKING:", "raise AssertionError", "raise NotImplementedError", '''class '.*\bProtocol\)':''', - # ''''@(abc\.)?'abstractmethod''', +#''''@(abc\.)?'abstractmethod''', ] fail_under = 40 show_missing = true @@ -166,4 +168,5 @@ omit = [ "setup.py", "elastica/systems/analytical.py", "elastica/experimental/*", + "elastica/**/protocol.py", ] From cba8e24d0bf22cd4ac43d9d482918b983c1ea50c Mon Sep 17 00:00:00 2001 From: Arman Tekinalp <53585636+armantekinalp@users.noreply.github.com> Date: Tue, 30 Apr 2024 11:08:05 -0500 Subject: [PATCH 019/134] Update elastica/timestepper/symplectic_steppers.py --- elastica/timestepper/symplectic_steppers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index 783b34cc6..645f9f256 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -45,7 +45,7 @@ def do_step( dyn_step(TimeStepper, System, time, dt) # Peel the last kinematic step and prefactor alone - last_kin_prefactor: Callable[..., np.floating] = _steps_and_prefactors[-1][0] + last_kin_prefactor: OperatorType = _steps_and_prefactors[-1][0] last_kin_step = _steps_and_prefactors[-1][1] last_kin_step(TimeStepper, System, time, dt) From df474ebe311bf47dbb0cc486a096c62979b886d2 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Tue, 30 Apr 2024 11:46:57 -0500 Subject: [PATCH 020/134] formatting: remove any unused imports --- elastica/systems/protocol.py | 2 +- elastica/timestepper/explicit_steppers.py | 2 +- elastica/timestepper/symplectic_steppers.py | 2 +- elastica/timestepper/tag.py | 2 +- elastica/typing.py | 4 +--- 5 files changed, 5 insertions(+), 7 deletions(-) diff --git a/elastica/systems/protocol.py b/elastica/systems/protocol.py index a9462536a..c2864f608 100644 --- a/elastica/systems/protocol.py +++ b/elastica/systems/protocol.py @@ -1,6 +1,6 @@ __doc__ = """Base class for elastica system""" -from typing import Protocol, Any +from typing import Protocol from elastica.typing import StateType from elastica.rod.data_structures import _KinematicState, _DynamicState diff --git a/elastica/timestepper/explicit_steppers.py b/elastica/timestepper/explicit_steppers.py index ca21f8dea..11ac06e97 100644 --- a/elastica/timestepper/explicit_steppers.py +++ b/elastica/timestepper/explicit_steppers.py @@ -1,6 +1,6 @@ __doc__ = """Explicit timesteppers and concepts""" -from typing import Any, Tuple +from typing import Tuple import numpy as np from copy import copy diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index 645f9f256..ea28aca6f 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -1,6 +1,6 @@ __doc__ = """Symplectic time steppers and concepts for integrating the kinematic and dynamic equations of rod-like objects. """ -from typing import Tuple, Callable, Any, List +from typing import Any, Callable, List from elastica.typing import ( SystemCollectionType, diff --git a/elastica/timestepper/tag.py b/elastica/timestepper/tag.py index 158f0c66b..85814e490 100644 --- a/elastica/timestepper/tag.py +++ b/elastica/timestepper/tag.py @@ -1,4 +1,4 @@ -from typing import Type, List, Callable +from typing import Callable, Type # TODO: Maybe move this for common utility diff --git a/elastica/typing.py b/elastica/typing.py index 6334bfd4c..b6bd654e6 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -5,9 +5,7 @@ from typing import TYPE_CHECKING from typing import Type, Union, Callable, Any -from typing import Protocol, TypeAlias - -import numpy as np +from typing import TypeAlias if TYPE_CHECKING: From f95925d843cd42aee2fc2157bfd65b59c3f31d2d Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Tue, 30 Apr 2024 12:07:13 -0500 Subject: [PATCH 021/134] Create py.typed --- elastica/py.typed | 1 + 1 file changed, 1 insertion(+) create mode 100644 elastica/py.typed diff --git a/elastica/py.typed b/elastica/py.typed new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/elastica/py.typed @@ -0,0 +1 @@ + From 71a566c3fc7236717f25cb01f633a70aa956f6ea Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Wed, 1 May 2024 10:30:52 -0500 Subject: [PATCH 022/134] refactor: Stateful -> Symplectic --- elastica/py.typed | 1 - elastica/timestepper/__init__.py | 37 +++++++++++++------------------ elastica/timestepper/protocol.py | 38 +++++++++++++------------------- elastica/typing.py | 19 ++++++++++------ 4 files changed, 42 insertions(+), 53 deletions(-) diff --git a/elastica/py.typed b/elastica/py.typed index 8b1378917..e69de29bb 100644 --- a/elastica/py.typed +++ b/elastica/py.typed @@ -1 +0,0 @@ - diff --git a/elastica/timestepper/__init__.py b/elastica/timestepper/__init__.py index 66f8190fe..10fed9d96 100644 --- a/elastica/timestepper/__init__.py +++ b/elastica/timestepper/__init__.py @@ -1,7 +1,7 @@ __doc__ = """Timestepping utilities to be used with Rod and RigidBody classes""" -from typing import Tuple, List, Callable, Type -from elastica.typing import SystemType +from typing import Tuple, List, Callable, Type, Any +from elastica.typing import SystemType, SystemCollectionType, SteppersOperatorsType import numpy as np from tqdm import tqdm @@ -12,30 +12,24 @@ from .explicit_steppers import RungeKutta4, EulerForward from .tag import SymplecticStepperTag, ExplicitStepperTag -from .protocol import StepperProtocol, StatefulStepperProtocol -from .protocol import MethodCollectorProtocol +from .protocol import StepperProtocol, SymplecticStepperProtocol # TODO: Both extend_stepper_interface and integrate should be in separate file. # __init__ is probably not an ideal place to have these scripts. def extend_stepper_interface( - Stepper: StepperProtocol, System: SystemType + Stepper: StepperProtocol, System: SystemType | SystemCollectionType ) -> Tuple[Callable, Tuple[Callable]]: - # StepperMethodCollector: Type[MethodCollectorProtocol] - # SystemStepper: Type[StepperProtocol] - if isinstance(Stepper.Tag, SymplecticStepperTag): + # SystemStepper: StepperProtocol + if Stepper.Tag == SymplecticStepperTag: from elastica.timestepper.symplectic_steppers import ( - _SystemInstanceStepper, - _SystemCollectionStepper, SymplecticStepperMethods, ) StepperMethodCollector = SymplecticStepperMethods - elif isinstance(Stepper.Tag, ExplicitStepperTag): # type: ignore[no-redef] + elif Stepper.Tag == ExplicitStepperTag: # type: ignore[no-redef] from elastica.timestepper.explicit_steppers import ( - _SystemInstanceStepper, - _SystemCollectionStepper, ExplicitStepperMethods, ) @@ -48,19 +42,18 @@ def extend_stepper_interface( ) # Check if system is a "collection" of smaller systems - if is_system_a_collection(System): - SystemStepper = _SystemCollectionStepper - else: - SystemStepper = _SystemInstanceStepper + assert is_system_a_collection(System) - stepper_methods: Tuple[Callable] = StepperMethodCollector(Stepper).step_methods() - do_step_method: Callable = SystemStepper.do_step + stepper_methods: SteppersOperatorsType = StepperMethodCollector( + Stepper + ).step_methods() + do_step_method: Callable = stepper_methods.do_step return do_step_method, stepper_methods def integrate( - StatefulStepper: StatefulStepperProtocol, - System: SystemType, + StatefulStepper: SymplecticStepperProtocol, + System: SystemType | SystemCollectionType, final_time: float, n_steps: int = 1000, restart_time: float = 0.0, @@ -70,7 +63,7 @@ def integrate( Parameters ---------- - StatefulStepper : StatefulStepperProtocol + StatefulStepper : SymplecticStepperProtocol Stepper algorithm to use. System : SystemType The elastica-system to simulate. diff --git a/elastica/timestepper/protocol.py b/elastica/timestepper/protocol.py index 63a9a52d6..27086a1ea 100644 --- a/elastica/timestepper/protocol.py +++ b/elastica/timestepper/protocol.py @@ -1,8 +1,9 @@ __doc__ = "Time stepper interface" -from typing import Protocol, Tuple, Callable, Type, Any +from typing import Protocol, Callable, Literal, ClassVar -from elastica.typing import SystemType +from elastica.typing import SystemType, SteppersOperatorsType, OperatorType +from .tag import StepperTags import numpy as np @@ -10,37 +11,28 @@ class StepperProtocol(Protocol): """Protocol for all time-steppers""" - def do_step(self, *args: Any, **kwargs: Any) -> float: ... + Tag: StepperTags @property - def Tag(self) -> Type: ... + def n_stages(self) -> int: ... + def step_methods(self) -> SteppersOperatorsType: ... -class StatefulStepperProtocol(StepperProtocol, Protocol): - """ - Stateful explicit, symplectic stepper wrapper. - """ - # For stateful steppes, bind memory to self - def do_step(self, System: SystemType, time: np.floating, dt: np.floating) -> float: - """ - Perform one time step of the simulation. - Return the new time. - """ - ... +class SymplecticStepperProtocol(StepperProtocol, Protocol): + """symplectic stepper protocol.""" - @property - def n_stages(self) -> int: ... + def get_steps(self) -> list[OperatorType]: ... + + def get_prefactors(self) -> list[OperatorType]: ... -class MethodCollectorProtocol(Protocol): - """ - Protocol for collecting stepper methods. - """ +class ExplicitStepperProtocol(StepperProtocol, Protocol): + """symplectic stepper protocol.""" - def __init__(self, timestepper_instance: StepperProtocol): ... + def get_steps(self) -> list[OperatorType]: ... - def step_methods(self) -> Tuple[Callable]: ... + def get_prefactors(self) -> list[OperatorType]: ... class MemoryProtocol(Protocol): diff --git a/elastica/typing.py b/elastica/typing.py index b6bd654e6..7c3df8702 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -18,7 +18,11 @@ from .rod.data_structures import _State as State from .systems.protocol import SymplecticSystemProtocol, ExplicitSystemProtocol - from .timestepper.protocol import StatefulStepperProtocol, MemoryProtocol + from .timestepper.protocol import ( + StepperProtocol, + SymplecticStepperProtocol, + MemoryProtocol, + ) else: RodBase = None RigidBodyBase = None @@ -28,7 +32,9 @@ State = "State" SymplecticSystemProtocol = None ExplicitSystemProtocol = None - StatefulStepperProtocol = None + + StepperProtocol = None + SymplecticStepperProtocol = None MemoryProtocol = None @@ -41,10 +47,10 @@ # NoOpt stepper # Symplectic stepper # StepOperatorType = Callable[ -# [StatefulStepperProtocol, SymplecticSystemProtocol, np.floating, np.floating], None +# [SymplecticStepperProtocol, SymplecticSystemProtocol, np.floating, np.floating], None # ] # PrefactorOperatorType = Callable[ -# [StatefulStepperProtocol, np.floating], np.floating +# [SymplecticStepperProtocol, np.floating], np.floating # ] OperatorType: TypeAlias = Callable[ Any, Any @@ -54,7 +60,7 @@ # Explicit stepper # ExplicitStageOperatorType = Callable[ # [ -# StatefulStepperProtocol, +# SymplecticStepperProtocol, # ExplicitSystemProtocol, # MemoryProtocol, # np.floating, @@ -64,7 +70,7 @@ # ] # ExplicitUpdateOperatorType = Callable[ # [ -# StatefulStepperProtocol, +# SymplecticStepperProtocol, # ExplicitSystemProtocol, # MemoryProtocol, # np.floating, @@ -72,7 +78,6 @@ # ], # np.floating, # ] -ExplicitOperatorsType: TypeAlias = tuple[tuple[OperatorType, ...], ...] RodType: TypeAlias = Type[RodBase] SystemCollectionType: TypeAlias = BaseSystemCollection From c033bd7dc30a92a458741eb56455b8b16938dc5d Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Wed, 1 May 2024 10:34:25 -0500 Subject: [PATCH 023/134] update: use string-literal taging instead of class-typing --- elastica/timestepper/explicit_steppers.py | 24 +-- elastica/timestepper/symplectic_steppers.py | 198 ++++++++------------ elastica/timestepper/tag.py | 33 +--- 3 files changed, 100 insertions(+), 155 deletions(-) diff --git a/elastica/timestepper/explicit_steppers.py b/elastica/timestepper/explicit_steppers.py index 11ac06e97..271256baf 100644 --- a/elastica/timestepper/explicit_steppers.py +++ b/elastica/timestepper/explicit_steppers.py @@ -8,12 +8,12 @@ from elastica.typing import ( SystemCollectionType, OperatorType, - ExplicitOperatorsType, + SteppersOperatorsType, StateType, ) from elastica.systems.protocol import ExplicitSystemProtocol -from .tag import tag, ExplicitStepperTag -from .protocol import StatefulStepperProtocol, MemoryProtocol +from .tag import StepperTags, ExplicitStepperTag +from .protocol import ExplicitStepperProtocol, MemoryProtocol """ @@ -69,8 +69,8 @@ class _SystemInstanceStepper: # # noinspection PyUnresolvedReferences @staticmethod def do_step( - TimeStepper: StatefulStepperProtocol, - _stages_and_updates: ExplicitOperatorsType, + TimeStepper: ExplicitStepperProtocol, + _stages_and_updates: SteppersOperatorsType, System: ExplicitSystemProtocol, Memory: MemoryProtocol, time: np.floating, @@ -86,8 +86,8 @@ class _SystemCollectionStepper: # # noinspection PyUnresolvedReferences @staticmethod def do_step( - TimeStepper: StatefulStepperProtocol, - _stages_and_updates: ExplicitOperatorsType, + TimeStepper: ExplicitStepperProtocol, + _stages_and_updates: SteppersOperatorsType, SystemCollection: SystemCollectionType, MemoryCollection: Tuple[MemoryProtocol, ...], time: np.floating, @@ -111,7 +111,7 @@ class ExplicitStepperMethods: Can also be used as a mixin with optional cls argument below """ - def __init__(self, timestepper_instance: StatefulStepperProtocol): + def __init__(self, timestepper_instance: ExplicitStepperProtocol): take_methods_from = timestepper_instance __stages: list[OperatorType] = [ v @@ -134,7 +134,7 @@ def __init__(self, timestepper_instance: StatefulStepperProtocol): self._stages_and_updates = tuple(zip(__stages, __updates)) - def step_methods(self) -> ExplicitOperatorsType: + def step_methods(self) -> SteppersOperatorsType: return self._stages_and_updates @property @@ -147,12 +147,13 @@ def __init__(self, initial_state: StateType) -> None: self.initial_state = initial_state -@tag(ExplicitStepperTag) class EulerForward: """ Classical Euler Forward stepper. Stateless, coordinates operations only. """ + Tag: StepperTags = ExplicitStepperTag + def _first_stage( self, System: ExplicitSystemProtocol, @@ -196,13 +197,14 @@ def __init__( self.k_4 = k_4 -@tag(ExplicitStepperTag) class RungeKutta4: """ Stateless runge-kutta4. coordinates operations only, memory needs to be externally managed and allocated. """ + Tag: StepperTags = ExplicitStepperTag + # These methods should be static, but because we need to enable automatic # discovery in ExplicitStepper, these are bound to the RungeKutta4 class # For automatic discovery, the order of declaring stages here is very important diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index ea28aca6f..7b6a34bb4 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -1,6 +1,8 @@ __doc__ = """Symplectic time steppers and concepts for integrating the kinematic and dynamic equations of rod-like objects. """ -from typing import Any, Callable, List +from typing import Callable, Any + +from itertools import zip_longest from elastica.typing import ( SystemCollectionType, @@ -17,8 +19,8 @@ overload_operator_dynamic_numba, ) from elastica.systems.protocol import SymplecticSystemProtocol -from .protocol import StatefulStepperProtocol -from .tag import tag, SymplecticStepperTag +from .protocol import SymplecticStepperProtocol +from .tag import StepperTags, SymplecticStepperTag """ Developer Note @@ -29,38 +31,53 @@ """ -class _SystemInstanceStepper: - @staticmethod - def do_step( - TimeStepper: StatefulStepperProtocol, - _steps_and_prefactors: SteppersOperatorsType, - System: SymplecticSystemProtocol, - time: np.floating, - dt: np.floating, - ) -> np.floating: - for kin_prefactor, kin_step, dyn_step in _steps_and_prefactors[:-1]: - kin_step(TimeStepper, System, time, dt) - time += kin_prefactor(TimeStepper, dt) - System.update_internal_forces_and_torques(time) - dyn_step(TimeStepper, System, time, dt) +class SymplecticStepperMixin: + def step_methods(self) -> SteppersOperatorsType: + # Let the total number of steps for the Symplectic method + # be (2*n + 1) (for time-symmetry). What we do is collect + # the first n + 1 entries down in _steps and _prefac below, and then + # reverse and append it to itself. + _steps: list[OperatorType] = self.get_steps() + _steps = _steps + _steps[-2::-1] - # Peel the last kinematic step and prefactor alone - last_kin_prefactor: OperatorType = _steps_and_prefactors[-1][0] - last_kin_step = _steps_and_prefactors[-1][1] + # Prefac here is necessary because the linear-exponential integrator + # needs only the prefactor and not the dt. + _prefactors: list[OperatorType] = self.get_prefactors() + _prefactors = _prefactors + _prefactors[-2::-1] + assert len(_steps) == 2 * len(_prefactors) - 1 - last_kin_step(TimeStepper, System, time, dt) - return time + last_kin_prefactor(TimeStepper, dt) + # Separate the kinematic and dynamic steps + _kinematic_steps: list[OperatorType] = _steps[::2] + _dynamic_steps: list[OperatorType] = _steps[1::2] + def no_operation(*args: Any) -> None: + pass -class _SystemCollectionStepper: - """ - Symplectic stepper collection class - """ + return tuple( + zip_longest( + _prefactors, + _kinematic_steps, + _dynamic_steps, + fillvalue=no_operation, + ) + ) + @property + def n_stages(self) -> int: + return len(self.get_prefactors()) + + def step( + self, SystemCollection: SystemCollectionType, time: np.floating, dt: np.floating + ): + steps_and_prefactors = self.step_methods() + return self.do_step(self, steps_and_prefactors, SystemCollection, time, dt) + + # TODO: Merge with .step method in the future. + # DEPRECATED: Use .step instead. @staticmethod def do_step( - TimeStepper: StatefulStepperProtocol, - _steps_and_prefactors: SteppersOperatorsType, + TimeStepper: SymplecticStepperProtocol, + steps_and_prefactors: SteppersOperatorsType, SystemCollection: SystemCollectionType, time: np.floating, dt: np.floating, @@ -68,17 +85,13 @@ def do_step( """ Function for doing symplectic stepper over the user defined rods (system). - Parameters - ---------- - SystemCollection: SystemCollectionType - time: float - dt: float - Returns ------- + time: float + The time after the integration step. """ - for kin_prefactor, kin_step, dyn_step in _steps_and_prefactors[:-1]: + for kin_prefactor, kin_step, dyn_step in steps_and_prefactors[:-1]: for system in SystemCollection._memory_blocks: kin_step(TimeStepper, system, time, dt) @@ -103,8 +116,8 @@ def do_step( SystemCollection.constrain_rates(time) # Peel the last kinematic step and prefactor alone - last_kin_prefactor = _steps_and_prefactors[-1][0] - last_kin_step = _steps_and_prefactors[-1][1] + last_kin_prefactor = steps_and_prefactors[-1][0] + last_kin_step = steps_and_prefactors[-1][1] for system in SystemCollection._memory_blocks: last_kin_step(TimeStepper, system, time, dt) @@ -121,87 +134,25 @@ def do_step( return time -class SymplecticStepperMethods: - def __init__(self, timestepper_instance: StatefulStepperProtocol): - take_methods_from = timestepper_instance - # Let the total number of steps for the Symplectic method - # be (2*n + 1) (for time-symmetry). What we do is collect - # the first n + 1 entries down in _steps and _prefac below, and then - # reverse and append it to itself. - self._steps: List[OperatorType] = [ - v - for (k, v) in take_methods_from.__class__.__dict__.items() - if k.endswith("step") - ] - # Prefac here is necessary because the linear-exponential integrator - # needs only the prefactor and not the dt. - self._prefactors: List[OperatorType] = [ - v - for (k, v) in take_methods_from.__class__.__dict__.items() - if k.endswith("prefactor") - ] - - def mirror(in_list: List[Callable]) -> None: - """Mirrors an input list ignoring the last element - If steps = [A, B, C] - then this call makes it [A, B, C, B, A] - - Parameters - ---------- - in_list : input list to be mirrored, modified in-place - - """ - # syntax is very ugly - if len(in_list) > 1: - in_list.extend(in_list[-2::-1]) - elif in_list: - in_list.append(in_list[0]) - - mirror(self._steps) - mirror(self._prefactors) - - assert ( - len(self._steps) == 2 * len(self._prefactors) - 1 - ), "Size mismatch in the number of steps and prefactors provided for a Symplectic Stepper!" - - self._kinematic_steps: List[OperatorType] = self._steps[::2] - self._dynamic_steps: List[OperatorType] = self._steps[1::2] - - # Avoid this check for MockClasses - if len(self._kinematic_steps) > 0: - assert ( - len(self._kinematic_steps) == len(self._dynamic_steps) + 1 - ), "Size mismatch in the number of kinematic and dynamic steps provided for a Symplectic Stepper!" - - from itertools import zip_longest - - def NoOp(*args: Any) -> None: - pass - - self._steps_and_prefactors: SteppersOperatorsType = tuple( - zip_longest( - self._prefactors, - self._kinematic_steps, - self._dynamic_steps, - fillvalue=NoOp, - ) - ) - - def step_methods(self) -> SteppersOperatorsType: - return self._steps_and_prefactors - - @property - def n_stages(self) -> int: - return len(self._steps_and_prefactors) - - -@tag(SymplecticStepperTag) -class PositionVerlet: +class PositionVerlet(SymplecticStepperMixin): """ Position Verlet symplectic time stepper class, which includes methods for second-order position Verlet. """ + Tag: StepperTags = SymplecticStepperTag + + def get_steps(self) -> list[OperatorType]: + return [ + self._first_kinematic_step, + self._first_dynamic_step, + ] + + def get_prefactors(self) -> list[OperatorType]: + return [ + self._first_prefactor, + ] + def _first_prefactor(self, dt: np.floating) -> np.floating: return 0.5 * dt @@ -227,14 +178,15 @@ def _first_dynamic_step( ) -@tag(SymplecticStepperTag) -class PEFRL: +class PEFRL(SymplecticStepperMixin): """ Position Extended Forest-Ruth Like Algorithm of I.M. Omelyan, I.M. Mryglod and R. Folk, Computer Physics Communications 146, 188 (2002), http://arxiv.org/abs/cond-mat/0110585 """ + Tag: StepperTags = SymplecticStepperTag + # xi and chi are confusing, but be careful! ξ: np.float64 = np.float64(0.1786178958448091e0) # ξ λ: np.float64 = -np.float64(0.2123418310626054e0) # λ @@ -244,6 +196,22 @@ class PEFRL: lambda_dash_coeff: np.float64 = 0.5 * (1.0 - 2.0 * λ) xi_chi_dash_coeff: np.float64 = 1.0 - 2.0 * (ξ + χ) + def get_steps(self) -> list[OperatorType]: + return [ + self._first_kinematic_step, + self._first_dynamic_step, + self._second_kinematic_step, + self._second_dynamic_step, + self._third_kinematic_step, + ] + + def get_prefactors(self) -> list[OperatorType]: + return [ + self._first_prefactor, + self._second_prefactor, + self._third_prefactor, + ] + def _first_kinematic_prefactor(self, dt: np.floating) -> np.floating: return self.ξ * dt diff --git a/elastica/timestepper/tag.py b/elastica/timestepper/tag.py index 85814e490..14a852832 100644 --- a/elastica/timestepper/tag.py +++ b/elastica/timestepper/tag.py @@ -1,31 +1,6 @@ -from typing import Callable, Type +from typing import Literal, TypeAlias, Final -# TODO: Maybe move this for common utility -def tag(Tag: Type) -> Callable[[Type], Type]: - """ - Tag a class with arbitrary type-class - - example: - class ATag: ... - - @tag(ATag) - class A1: - ... - - assert isinstance(A1.tag, ATag) - """ - - def wrapper(cls: Type) -> Type: - cls.Tag = Tag() - return cls - - return wrapper - - -class SymplecticStepperTag: - pass - - -class ExplicitStepperTag: - pass +ExplicitStepperTag: str = "ExplicitStepper" +SymplecticStepperTag: str = "SymplecticStepper" +StepperTags: TypeAlias = Literal["SymplecticStepper", "ExplicitStepper"] From 0b47d9a8cdee0cf81d02f14e48a137af2561d466 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Fri, 3 May 2024 22:15:53 +0000 Subject: [PATCH 024/134] build(deps): bump tqdm from 4.66.2 to 4.66.3 Bumps [tqdm](https://github.com/tqdm/tqdm) from 4.66.2 to 4.66.3. - [Release notes](https://github.com/tqdm/tqdm/releases) - [Commits](https://github.com/tqdm/tqdm/compare/v4.66.2...v4.66.3) --- updated-dependencies: - dependency-name: tqdm dependency-type: direct:production ... Signed-off-by: dependabot[bot] --- poetry.lock | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/poetry.lock b/poetry.lock index a45d2d19b..4c0f39226 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.5.1 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand. [[package]] name = "accessible-pygments" @@ -1485,6 +1485,7 @@ files = [ {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:69b023b2b4daa7548bcfbd4aa3da05b3a74b772db9e23b982788168117739938"}, {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:81e0b275a9ecc9c0c0c07b4b90ba548307583c125f54d5b6946cfee6360c733d"}, {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba336e390cd8e4d1739f42dfe9bb83a3cc2e80f567d8805e11b46f4a943f5515"}, + {file = "PyYAML-6.0.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:326c013efe8048858a6d312ddd31d56e468118ad4cdeda36c719bf5bb6192290"}, {file = "PyYAML-6.0.1-cp310-cp310-win32.whl", hash = "sha256:bd4af7373a854424dabd882decdc5579653d7868b8fb26dc7d0e99f823aa5924"}, {file = "PyYAML-6.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:fd1592b3fdf65fff2ad0004b5e363300ef59ced41c2e6b3a99d4089fa8c5435d"}, {file = "PyYAML-6.0.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:6965a7bc3cf88e5a1c3bd2e0b5c22f8d677dc88a455344035f03399034eb3007"}, @@ -1492,8 +1493,16 @@ files = [ {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:42f8152b8dbc4fe7d96729ec2b99c7097d656dc1213a3229ca5383f973a5ed6d"}, {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:062582fca9fabdd2c8b54a3ef1c978d786e0f6b3a1510e0ac93ef59e0ddae2bc"}, {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d2b04aac4d386b172d5b9692e2d2da8de7bfb6c387fa4f801fbf6fb2e6ba4673"}, + {file = "PyYAML-6.0.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e7d73685e87afe9f3b36c799222440d6cf362062f78be1013661b00c5c6f678b"}, {file = "PyYAML-6.0.1-cp311-cp311-win32.whl", hash = "sha256:1635fd110e8d85d55237ab316b5b011de701ea0f29d07611174a1b42f1444741"}, {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"}, + {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"}, + {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"}, + {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"}, + {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"}, + {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"}, + {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"}, + {file = "PyYAML-6.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:0d3304d8c0adc42be59c5f8a4d9e3d7379e6955ad754aa9d6ab7a398b59dd1df"}, {file = "PyYAML-6.0.1-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:50550eb667afee136e9a77d6dc71ae76a44df8b3e51e41b77f6de2932bfe0f47"}, {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1fe35611261b29bd1de0070f0b2f47cb6ff71fa6595c077e42bd0c419fa27b98"}, {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:704219a11b772aea0d8ecd7058d0082713c3562b4e271b849ad7dc4a5c90c13c"}, @@ -1510,6 +1519,7 @@ files = [ {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a0cd17c15d3bb3fa06978b4e8958dcdc6e0174ccea823003a106c7d4d7899ac5"}, {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:28c119d996beec18c05208a8bd78cbe4007878c6dd15091efb73a30e90539696"}, {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7e07cbde391ba96ab58e532ff4803f79c4129397514e1413a7dc761ccd755735"}, + {file = "PyYAML-6.0.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:49a183be227561de579b4a36efbb21b3eab9651dd81b1858589f796549873dd6"}, {file = "PyYAML-6.0.1-cp38-cp38-win32.whl", hash = "sha256:184c5108a2aca3c5b3d3bf9395d50893a7ab82a38004c8f61c258d4428e80206"}, {file = "PyYAML-6.0.1-cp38-cp38-win_amd64.whl", hash = "sha256:1e2722cc9fbb45d9b87631ac70924c11d3a401b2d7f410cc0e3bbf249f2dca62"}, {file = "PyYAML-6.0.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:9eb6caa9a297fc2c2fb8862bc5370d0303ddba53ba97e71f08023b6cd73d16a8"}, @@ -1517,6 +1527,7 @@ files = [ {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5773183b6446b2c99bb77e77595dd486303b4faab2b086e7b17bc6bef28865f6"}, {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b786eecbdf8499b9ca1d697215862083bd6d2a99965554781d0d8d1ad31e13a0"}, {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bc1bf2925a1ecd43da378f4db9e4f799775d6367bdb94671027b73b393a7c42c"}, + {file = "PyYAML-6.0.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:04ac92ad1925b2cff1db0cfebffb6ffc43457495c9b3c39d3fcae417d7125dc5"}, {file = "PyYAML-6.0.1-cp39-cp39-win32.whl", hash = "sha256:faca3bdcf85b2fc05d06ff3fbc1f83e1391b3e724afa3feba7d13eeab355484c"}, {file = "PyYAML-6.0.1-cp39-cp39-win_amd64.whl", hash = "sha256:510c9deebc5c0225e8c96813043e62b680ba2f9c50a08d3724c7f28a747d1486"}, {file = "PyYAML-6.0.1.tar.gz", hash = "sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43"}, @@ -1853,13 +1864,13 @@ files = [ [[package]] name = "tqdm" -version = "4.66.2" +version = "4.66.3" description = "Fast, Extensible Progress Meter" optional = false python-versions = ">=3.7" files = [ - {file = "tqdm-4.66.2-py3-none-any.whl", hash = "sha256:1ee4f8a893eb9bef51c6e35730cebf234d5d0b6bd112b0271e10ed7c24a02bd9"}, - {file = "tqdm-4.66.2.tar.gz", hash = "sha256:6cd52cdf0fef0e0f543299cfc96fec90d7b8a7e88745f411ec33eb44d5ed3531"}, + {file = "tqdm-4.66.3-py3-none-any.whl", hash = "sha256:4f41d54107ff9a223dca80b53efe4fb654c67efaba7f47bada3ee9d50e05bd53"}, + {file = "tqdm-4.66.3.tar.gz", hash = "sha256:23097a41eba115ba99ecae40d06444c15d1c0c698d527a01c6c8bd1c5d0647e5"}, ] [package.dependencies] From edcc4f9f765af4ddc4b7980eaa981c83a544149c Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Wed, 1 May 2024 10:39:38 -0500 Subject: [PATCH 025/134] redesign integration flow: mypy checked --- elastica/timestepper/__init__.py | 61 +++---- elastica/timestepper/explicit_steppers.py | 169 ++++++++++---------- elastica/timestepper/protocol.py | 29 ++-- elastica/timestepper/symplectic_steppers.py | 19 ++- elastica/timestepper/tag.py | 4 +- examples/CatenaryCase/catenary.py | 13 +- tests/test_math/test_timestepper.py | 53 ++---- 7 files changed, 156 insertions(+), 192 deletions(-) diff --git a/elastica/timestepper/__init__.py b/elastica/timestepper/__init__.py index 10fed9d96..d385edc28 100644 --- a/elastica/timestepper/__init__.py +++ b/elastica/timestepper/__init__.py @@ -15,57 +15,39 @@ from .protocol import StepperProtocol, SymplecticStepperProtocol -# TODO: Both extend_stepper_interface and integrate should be in separate file. -# __init__ is probably not an ideal place to have these scripts. +# Deprecated: Remove in the future version +# Many script still uses this method to control timestep. Keep it for backward compatibility def extend_stepper_interface( Stepper: StepperProtocol, System: SystemType | SystemCollectionType -) -> Tuple[Callable, Tuple[Callable]]: - - # SystemStepper: StepperProtocol - if Stepper.Tag == SymplecticStepperTag: - from elastica.timestepper.symplectic_steppers import ( - SymplecticStepperMethods, - ) - - StepperMethodCollector = SymplecticStepperMethods - elif Stepper.Tag == ExplicitStepperTag: # type: ignore[no-redef] - from elastica.timestepper.explicit_steppers import ( - ExplicitStepperMethods, - ) - - StepperMethodCollector = ExplicitStepperMethods - else: - raise NotImplementedError( - "Only explicit and symplectic steppers are supported, given stepper is {}".format( - Stepper.__class__.__name__ - ) - ) - +) -> Tuple[ + Callable[ + [StepperProtocol, SystemCollectionType, np.floating, np.floating], np.floating + ], + SteppersOperatorsType, +]: # Check if system is a "collection" of smaller systems assert is_system_a_collection(System) - stepper_methods: SteppersOperatorsType = StepperMethodCollector( - Stepper - ).step_methods() - do_step_method: Callable = stepper_methods.do_step + stepper_methods: SteppersOperatorsType = Stepper.step_methods() + do_step_method: Callable = Stepper.do_step # type: ignore[attr-defined] return do_step_method, stepper_methods def integrate( - StatefulStepper: SymplecticStepperProtocol, - System: SystemType | SystemCollectionType, + Stepper: StepperProtocol, + SystemCollection: SystemCollectionType, final_time: float, n_steps: int = 1000, restart_time: float = 0.0, progress_bar: bool = True, -) -> float: +) -> np.floating: """ Parameters ---------- - StatefulStepper : SymplecticStepperProtocol + Stepper : StepperProtocol Stepper algorithm to use. - System : SystemType + SystemCollection : SystemType The elastica-system to simulate. final_time : float Total simulation time. The timestep is determined by final_time / n_steps. @@ -78,18 +60,13 @@ def integrate( """ assert final_time > 0.0, "Final time is negative!" assert n_steps > 0, "Number of integration steps is negative!" + assert is_system_a_collection(SystemCollection) - # Extend the stepper's interface after introspecting the properties - # of the system. If system is a collection of small systems (whose - # states cannot be aggregated), then stepper now loops over the system - # state - do_step, stages_and_updates = extend_stepper_interface(StatefulStepper, System) - - dt = np.float64(float(final_time) / n_steps) - time = restart_time + dt = np.float_(float(final_time) / n_steps) + time = np.float_(restart_time) for i in tqdm(range(n_steps), disable=(not progress_bar)): - time = do_step(StatefulStepper, stages_and_updates, System, time, dt) + time = Stepper.step(SystemCollection, time, dt) print("Final time of simulation is : ", time) return time diff --git a/elastica/timestepper/explicit_steppers.py b/elastica/timestepper/explicit_steppers.py index 271256baf..9b564e2e2 100644 --- a/elastica/timestepper/explicit_steppers.py +++ b/elastica/timestepper/explicit_steppers.py @@ -1,6 +1,6 @@ __doc__ = """Explicit timesteppers and concepts""" -from typing import Tuple +from typing import Type, Any import numpy as np from copy import copy @@ -65,35 +65,78 @@ """ -class _SystemInstanceStepper: - # # noinspection PyUnresolvedReferences - @staticmethod - def do_step( - TimeStepper: ExplicitStepperProtocol, - _stages_and_updates: SteppersOperatorsType, - System: ExplicitSystemProtocol, - Memory: MemoryProtocol, +class EulerForwardMemory: + def __init__(self, initial_state: StateType) -> None: + self.initial_state = initial_state + + +class RungeKutta4Memory: + """ + Stores all states of Rk within the time-stepper. Works as long as the states + are all one big numpy array, made possible by carefully using views. + + Convenience wrapper around Stateless that provides memory + """ + + def __init__( + self, + initial_state: StateType, + ) -> None: + self.initial_state = initial_state + self.k_1 = initial_state + self.k_2 = initial_state + self.k_3 = initial_state + self.k_4 = initial_state + + +class ExplicitStepperMixin: + """Base class for all explicit steppers + Can also be used as a mixin with optional cls argument below + """ + + def step_methods(self: ExplicitStepperProtocol) -> SteppersOperatorsType: + stages = self.get_stages() + updates = self.get_updates() + + assert len(stages) == len( + updates + ), "Number of stages and updates should be equal to one another" + return tuple(zip(stages, updates)) + + @property + def n_stages(self: ExplicitStepperProtocol) -> int: + return len(self.get_stages()) + + def step( + self: ExplicitStepperProtocol, + SystemCollection: SystemCollectionType, time: np.floating, dt: np.floating, ) -> np.floating: - for stage, update in _stages_and_updates: - stage(TimeStepper, System, Memory, time, dt) - time = update(TimeStepper, System, Memory, time, dt) - return time - + steps_and_prefactors = self.step_methods() + if isinstance( + self, EulerForward + ): # TODO: Cleanup - use depedency injection instead + Memory = EulerForwardMemory + elif isinstance(self, RungeKutta4): + Memory = RungeKutta4Memory # type: ignore[assignment] + else: + raise NotImplementedError(f"Memory class not defined for {self}") + memory_collection = tuple( + [Memory(initial_state=system.state) for system in SystemCollection] + ) + return ExplicitStepperMixin.do_step(self, steps_and_prefactors, SystemCollection, memory_collection, time, dt) # type: ignore[attr-defined] -class _SystemCollectionStepper: - # # noinspection PyUnresolvedReferences @staticmethod def do_step( TimeStepper: ExplicitStepperProtocol, - _stages_and_updates: SteppersOperatorsType, + stages_and_updates: SteppersOperatorsType, SystemCollection: SystemCollectionType, - MemoryCollection: Tuple[MemoryProtocol, ...], + MemoryCollection: Any, # TODO time: np.floating, dt: np.floating, ) -> np.floating: - for stage, update in _stages_and_updates: + for stage, update in stages_and_updates: SystemCollection.synchronize(time) for system, memory in zip(SystemCollection[:-1], MemoryCollection[:-1]): stage(TimeStepper, system, memory, time, dt) @@ -106,54 +149,19 @@ def do_step( return time -class ExplicitStepperMethods: - """Base class for all explicit steppers - Can also be used as a mixin with optional cls argument below - """ - - def __init__(self, timestepper_instance: ExplicitStepperProtocol): - take_methods_from = timestepper_instance - __stages: list[OperatorType] = [ - v - for (k, v) in take_methods_from.__class__.__dict__.items() - if k.endswith("stage") - ] - __updates: list[OperatorType] = [ - v - for (k, v) in take_methods_from.__class__.__dict__.items() - if k.endswith("update") - ] - - # Tuples are almost immutable - _n_stages: int = len(__stages) - _n_updates: int = len(__updates) - - assert ( - _n_stages == _n_updates - ), "Number of stages and updates should be equal to one another" - - self._stages_and_updates = tuple(zip(__stages, __updates)) - - def step_methods(self) -> SteppersOperatorsType: - return self._stages_and_updates - - @property - def n_stages(self) -> int: - return len(self._stages_and_updates) - - -class EulerForwardMemory: - def __init__(self, initial_state: StateType) -> None: - self.initial_state = initial_state - - -class EulerForward: +class EulerForward(ExplicitStepperMixin): """ Classical Euler Forward stepper. Stateless, coordinates operations only. """ Tag: StepperTags = ExplicitStepperTag + def get_stages(self) -> list[OperatorType]: + return [self._first_stage] + + def get_updates(self) -> list[OperatorType]: + return [self._first_update] + def _first_stage( self, System: ExplicitSystemProtocol, @@ -174,29 +182,6 @@ def _first_update( return time + dt -class RungeKutta4Memory: - """ - Stores all states of Rk within the time-stepper. Works as long as the states - are all one big numpy array, made possible by carefully using views. - - Convenience wrapper around Stateless that provides memory - """ - - def __init__( - self, - initial_state: StateType, - k_1: StateType, - k_2: StateType, - k_3: StateType, - k_4: StateType, - ) -> None: - self.initial_state = initial_state - self.k_1 = k_1 - self.k_2 = k_2 - self.k_3 = k_3 - self.k_4 = k_4 - - class RungeKutta4: """ Stateless runge-kutta4. coordinates operations only, memory needs @@ -205,6 +190,22 @@ class RungeKutta4: Tag: StepperTags = ExplicitStepperTag + def get_stages(self) -> list[OperatorType]: + return [ + self._first_stage, + self._second_stage, + self._third_stage, + self._fourth_stage, + ] + + def get_updates(self) -> list[OperatorType]: + return [ + self._first_update, + self._second_update, + self._third_update, + self._fourth_update, + ] + # These methods should be static, but because we need to enable automatic # discovery in ExplicitStepper, these are bound to the RungeKutta4 class # For automatic discovery, the order of declaring stages here is very important diff --git a/elastica/timestepper/protocol.py b/elastica/timestepper/protocol.py index 27086a1ea..f9101160b 100644 --- a/elastica/timestepper/protocol.py +++ b/elastica/timestepper/protocol.py @@ -1,8 +1,13 @@ __doc__ = "Time stepper interface" -from typing import Protocol, Callable, Literal, ClassVar - -from elastica.typing import SystemType, SteppersOperatorsType, OperatorType +from typing import Protocol, Callable, Literal, ClassVar, Type + +from elastica.typing import ( + SystemType, + SteppersOperatorsType, + OperatorType, + SystemCollectionType, +) from .tag import StepperTags import numpy as np @@ -18,6 +23,10 @@ def n_stages(self) -> int: ... def step_methods(self) -> SteppersOperatorsType: ... + def step( + self, SystemCollection: SystemCollectionType, time: np.floating, dt: np.floating + ) -> np.floating: ... + class SymplecticStepperProtocol(StepperProtocol, Protocol): """symplectic stepper protocol.""" @@ -27,17 +36,17 @@ def get_steps(self) -> list[OperatorType]: ... def get_prefactors(self) -> list[OperatorType]: ... -class ExplicitStepperProtocol(StepperProtocol, Protocol): - """symplectic stepper protocol.""" +class MemoryProtocol(Protocol): + @property + def initial_state(self) -> bool: ... - def get_steps(self) -> list[OperatorType]: ... - def get_prefactors(self) -> list[OperatorType]: ... +class ExplicitStepperProtocol(StepperProtocol, Protocol): + """symplectic stepper protocol.""" + def get_stages(self) -> list[OperatorType]: ... -class MemoryProtocol(Protocol): - @property - def initial_state(self) -> bool: ... + def get_updates(self) -> list[OperatorType]: ... # class _LinearExponentialIntegratorMixin: diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index 7b6a34bb4..f415680e2 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -32,7 +32,7 @@ class SymplecticStepperMixin: - def step_methods(self) -> SteppersOperatorsType: + def step_methods(self: SymplecticStepperProtocol) -> SteppersOperatorsType: # Let the total number of steps for the Symplectic method # be (2*n + 1) (for time-symmetry). What we do is collect # the first n + 1 entries down in _steps and _prefac below, and then @@ -63,14 +63,17 @@ def no_operation(*args: Any) -> None: ) @property - def n_stages(self) -> int: + def n_stages(self: SymplecticStepperProtocol) -> int: return len(self.get_prefactors()) def step( - self, SystemCollection: SystemCollectionType, time: np.floating, dt: np.floating - ): + self: SymplecticStepperProtocol, + SystemCollection: SystemCollectionType, + time: np.floating, + dt: np.floating, + ) -> np.floating: steps_and_prefactors = self.step_methods() - return self.do_step(self, steps_and_prefactors, SystemCollection, time, dt) + return SymplecticStepperMixin.do_step(self, steps_and_prefactors, SystemCollection, time, dt) # type: ignore[attr-defined] # TODO: Merge with .step method in the future. # DEPRECATED: Use .step instead. @@ -207,9 +210,9 @@ def get_steps(self) -> list[OperatorType]: def get_prefactors(self) -> list[OperatorType]: return [ - self._first_prefactor, - self._second_prefactor, - self._third_prefactor, + self._first_kinematic_prefactor, + self._second_kinematic_prefactor, + self._third_kinematic_prefactor, ] def _first_kinematic_prefactor(self, dt: np.floating) -> np.floating: diff --git a/elastica/timestepper/tag.py b/elastica/timestepper/tag.py index 14a852832..d23788750 100644 --- a/elastica/timestepper/tag.py +++ b/elastica/timestepper/tag.py @@ -1,6 +1,6 @@ from typing import Literal, TypeAlias, Final -ExplicitStepperTag: str = "ExplicitStepper" -SymplecticStepperTag: str = "SymplecticStepper" +ExplicitStepperTag: Final = "ExplicitStepper" +SymplecticStepperTag: Final = "SymplecticStepper" StepperTags: TypeAlias = Literal["SymplecticStepper", "ExplicitStepper"] diff --git a/examples/CatenaryCase/catenary.py b/examples/CatenaryCase/catenary.py index d6fd787b0..6a2d9e88e 100644 --- a/examples/CatenaryCase/catenary.py +++ b/examples/CatenaryCase/catenary.py @@ -1,5 +1,6 @@ import numpy as np from elastica import * +from elastica.typing import SystemType, SystemCollectionType, SymplecticStepperProtocol from post_processing import ( plot_video, @@ -11,7 +12,7 @@ class CatenarySimulator(BaseSystemCollection, Constraints, Forcing, Damping, Cal pass -catenary_sim = CatenarySimulator() +catenary_sim: SystemCollectionType = CatenarySimulator() final_time = 10 damping_constant = 0.3 time_step = 1e-4 @@ -37,7 +38,7 @@ class CatenarySimulator(BaseSystemCollection, Constraints, Forcing, Damping, Cal poisson_ratio = 0.5 shear_modulus = E / (poisson_ratio + 1.0) -base_rod = CosseratRod.straight_rod( +base_rod: SystemType = CosseratRod.straight_rod( n_elem, start, direction, @@ -76,11 +77,11 @@ class CatenaryCallBack(CallBackBaseClass): """ def __init__(self, step_skip: int, callback_params: dict): - CallBackBaseClass.__init__(self) + super().__init__() self.every = step_skip self.callback_params = callback_params - def make_callback(self, system, time, current_step: int): + def make_callback(self, system: SystemType, time: float, current_step: int) -> None: if current_step % self.every == 0: @@ -93,7 +94,7 @@ def make_callback(self, system, time, current_step: int): return -pp_list = defaultdict(list) +pp_list: dict = defaultdict(list) catenary_sim.collect_diagnostics(base_rod).using( CatenaryCallBack, step_skip=step_skip, callback_params=pp_list ) @@ -102,7 +103,7 @@ def make_callback(self, system, time, current_step: int): catenary_sim.finalize() -timestepper = PositionVerlet() +timestepper: SymplecticStepperProtocol = PositionVerlet() integrate(timestepper, catenary_sim, final_time, total_steps) position = np.array(pp_list["position"]) diff --git a/tests/test_math/test_timestepper.py b/tests/test_math/test_timestepper.py index 78c9cf605..98599a5aa 100644 --- a/tests/test_math/test_timestepper.py +++ b/tests/test_math/test_timestepper.py @@ -18,17 +18,14 @@ from elastica.timestepper.explicit_steppers import ( RungeKutta4, EulerForward, + ExplicitStepperMixin, ) - -# from elastica.timestepper.explicit_steppers import ( -# StatefulRungeKutta4, -# StatefulEulerForward, -# ) from elastica.timestepper.symplectic_steppers import ( PositionVerlet, PEFRL, + SymplecticStepperMixin, ) -from elastica.timestepper.tag import tag, SymplecticStepperTag, ExplicitStepperTag +from elastica.timestepper.tag import SymplecticStepperTag, ExplicitStepperTag from elastica.utils import Tolerance @@ -37,8 +34,8 @@ class TestExtendStepperInterface: """TODO add documentation""" - @tag(SymplecticStepperTag) - class MockSymplecticStepper: + class MockSymplecticStepper(SymplecticStepperMixin): + Tag = SymplecticStepperTag def _first_prefactor(self): pass @@ -49,8 +46,8 @@ def _first_kinematic_step(self): def _first_dynamic_step(self): pass - @tag(ExplicitStepperTag) - class MockExplicitStepper: + class MockExplicitStepper(ExplicitStepperMixin): + Tag = ExplicitStepperTag def _first_stage(self): pass @@ -58,20 +55,6 @@ def _first_stage(self): def _first_update(self): pass - from elastica.timestepper.symplectic_steppers import ( - _SystemInstanceStepper as symplectic_instance_stepper, - ) - from elastica.timestepper.symplectic_steppers import ( - _SystemCollectionStepper as symplectic_collection_stepper, - ) - - from elastica.timestepper.explicit_steppers import ( - _SystemInstanceStepper as explicit_instance_stepper, - ) - from elastica.timestepper.explicit_steppers import ( - _SystemCollectionStepper as explicit_collection_stepper, - ) - # We cannot call a stepper on a system until both the stepper # and system "see" one another (for performance reasons, mostly) # So before "seeing" the system, the stepper should not have @@ -79,16 +62,13 @@ def _first_update(self): # after "seeing" the system, via extend_stepper_interface @pytest.mark.parametrize( "stepper_and_interface", - [ - (MockSymplecticStepper, symplectic_instance_stepper), - (MockExplicitStepper, explicit_instance_stepper), - ], + [MockSymplecticStepper, MockExplicitStepper], ) def test_symplectic_stepper_interface_for_simple_systems( self, stepper_and_interface ): system = ScalarExponentialDecaySystem() - (stepper_cls, interface_cls) = stepper_and_interface + stepper_cls = stepper_and_interface stepper = stepper_cls() stepper_methods = None @@ -100,16 +80,13 @@ def test_symplectic_stepper_interface_for_simple_systems( @pytest.mark.parametrize( "stepper_and_interface", - [ - (MockSymplecticStepper, symplectic_collection_stepper), - (MockExplicitStepper, explicit_collection_stepper), - ], + [MockSymplecticStepper, MockExplicitStepper], ) def test_symplectic_stepper_interface_for_collective_systems( self, stepper_and_interface ): system = SymplecticUndampedHarmonicOscillatorCollectiveSystem() - (stepper_cls, interface_cls) = stepper_and_interface + stepper_cls = stepper_and_interface stepper = stepper_cls() stepper_methods = None @@ -122,16 +99,12 @@ def test_symplectic_stepper_interface_for_collective_systems( class MockBadStepper: Tag = int() # an arbitrary tag that doesn't mean anything - @pytest.mark.parametrize( - "stepper_and_interface", [(MockBadStepper, symplectic_collection_stepper)] - ) + @pytest.mark.parametrize("stepper_and_interface", [MockBadStepper]) def test_symplectic_stepper_throws_for_bad_stepper(self, stepper_and_interface): system = ScalarExponentialDecaySystem() - (stepper_cls, interface_cls) = stepper_and_interface + stepper_cls = stepper_and_interface stepper = stepper_cls() - assert interface_cls not in stepper.__class__.__bases__ - with pytest.raises(NotImplementedError) as excinfo: extend_stepper_interface(stepper, system) assert "steppers are supported" in str(excinfo.value) From 643fa7dfb61dc82da41757ef45210ca3fb5ae2c6 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Thu, 2 May 2024 10:41:52 -0500 Subject: [PATCH 026/134] update: save operator sequence without running at each timestep --- elastica/timestepper/explicit_steppers.py | 10 ++++++---- elastica/timestepper/protocol.py | 5 +++++ elastica/timestepper/symplectic_steppers.py | 10 ++++++---- 3 files changed, 17 insertions(+), 8 deletions(-) diff --git a/elastica/timestepper/explicit_steppers.py b/elastica/timestepper/explicit_steppers.py index 9b564e2e2..3064f00e1 100644 --- a/elastica/timestepper/explicit_steppers.py +++ b/elastica/timestepper/explicit_steppers.py @@ -1,6 +1,6 @@ __doc__ = """Explicit timesteppers and concepts""" -from typing import Type, Any +from typing import Type, Any, Final import numpy as np from copy import copy @@ -94,6 +94,9 @@ class ExplicitStepperMixin: Can also be used as a mixin with optional cls argument below """ + def __init__(self: ExplicitStepperProtocol): + self.steps_and_prefactors = self.step_methods() + def step_methods(self: ExplicitStepperProtocol) -> SteppersOperatorsType: stages = self.get_stages() updates = self.get_updates() @@ -105,7 +108,7 @@ def step_methods(self: ExplicitStepperProtocol) -> SteppersOperatorsType: @property def n_stages(self: ExplicitStepperProtocol) -> int: - return len(self.get_stages()) + return len(self.steps_and_prefactors) def step( self: ExplicitStepperProtocol, @@ -113,7 +116,6 @@ def step( time: np.floating, dt: np.floating, ) -> np.floating: - steps_and_prefactors = self.step_methods() if isinstance( self, EulerForward ): # TODO: Cleanup - use depedency injection instead @@ -125,7 +127,7 @@ def step( memory_collection = tuple( [Memory(initial_state=system.state) for system in SystemCollection] ) - return ExplicitStepperMixin.do_step(self, steps_and_prefactors, SystemCollection, memory_collection, time, dt) # type: ignore[attr-defined] + return ExplicitStepperMixin.do_step(self, self.steps_and_prefactors, SystemCollection, memory_collection, time, dt) # type: ignore[attr-defined] @staticmethod def do_step( diff --git a/elastica/timestepper/protocol.py b/elastica/timestepper/protocol.py index f9101160b..f923e62eb 100644 --- a/elastica/timestepper/protocol.py +++ b/elastica/timestepper/protocol.py @@ -18,9 +18,14 @@ class StepperProtocol(Protocol): Tag: StepperTags + def __init__(self) -> None: ... + @property def n_stages(self) -> int: ... + @property + def steps_and_prefactors(self) -> SteppersOperatorsType: ... + def step_methods(self) -> SteppersOperatorsType: ... def step( diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index f415680e2..a09124594 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -1,6 +1,6 @@ __doc__ = """Symplectic time steppers and concepts for integrating the kinematic and dynamic equations of rod-like objects. """ -from typing import Callable, Any +from typing import Callable, Any, Final from itertools import zip_longest @@ -32,6 +32,9 @@ class SymplecticStepperMixin: + def __init__(self: SymplecticStepperProtocol): + self.steps_and_prefactors: Final = self.step_methods() + def step_methods(self: SymplecticStepperProtocol) -> SteppersOperatorsType: # Let the total number of steps for the Symplectic method # be (2*n + 1) (for time-symmetry). What we do is collect @@ -64,7 +67,7 @@ def no_operation(*args: Any) -> None: @property def n_stages(self: SymplecticStepperProtocol) -> int: - return len(self.get_prefactors()) + return len(self.steps_and_prefactors) def step( self: SymplecticStepperProtocol, @@ -72,8 +75,7 @@ def step( time: np.floating, dt: np.floating, ) -> np.floating: - steps_and_prefactors = self.step_methods() - return SymplecticStepperMixin.do_step(self, steps_and_prefactors, SystemCollection, time, dt) # type: ignore[attr-defined] + return SymplecticStepperMixin.do_step(self, self.steps_and_prefactors, SystemCollection, time, dt) # type: ignore[attr-defined] # TODO: Merge with .step method in the future. # DEPRECATED: Use .step instead. From c5a9c5acee4c9ac7f618121cf790e7dad96f1e0f Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Thu, 2 May 2024 13:58:35 -0500 Subject: [PATCH 027/134] test: remove symple system stepping - no longer support --- tests/test_math/test_timestepper.py | 24 +++--------------------- 1 file changed, 3 insertions(+), 21 deletions(-) diff --git a/tests/test_math/test_timestepper.py b/tests/test_math/test_timestepper.py index 98599a5aa..3868013f9 100644 --- a/tests/test_math/test_timestepper.py +++ b/tests/test_math/test_timestepper.py @@ -61,32 +61,14 @@ def _first_update(self): # the interface (interface_cls). It should however have the interface # after "seeing" the system, via extend_stepper_interface @pytest.mark.parametrize( - "stepper_and_interface", - [MockSymplecticStepper, MockExplicitStepper], - ) - def test_symplectic_stepper_interface_for_simple_systems( - self, stepper_and_interface - ): - system = ScalarExponentialDecaySystem() - stepper_cls = stepper_and_interface - stepper = stepper_cls() - - stepper_methods = None - assert stepper_methods is None - - _, stepper_methods = extend_stepper_interface(stepper, system) - - assert stepper_methods - - @pytest.mark.parametrize( - "stepper_and_interface", + "stepper_module", [MockSymplecticStepper, MockExplicitStepper], ) def test_symplectic_stepper_interface_for_collective_systems( - self, stepper_and_interface + self, stepper_module ): system = SymplecticUndampedHarmonicOscillatorCollectiveSystem() - stepper_cls = stepper_and_interface + stepper_cls = stepper_module.steps_and_prefactors stepper = stepper_cls() stepper_methods = None From 3fb7dac7b8f6ecd7b2b4f4db04260ff729ed92b6 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Thu, 2 May 2024 14:31:25 -0500 Subject: [PATCH 028/134] fix: tuple is not supposed to be mutable --- elastica/modules/base_system.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index 34f42c7d4..4c3d46528 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -5,7 +5,8 @@ Basic coordinating for multiple, smaller systems that have an independently integrable interface (i.e. works with symplectic or explicit routines `timestepper.py`.) """ -from typing import Iterable, Callable, AnyStr +from typing import Iterable, Callable, AnyStr, Type +from elastica.typing import SystemType import numpy as np @@ -57,7 +58,7 @@ def __init__(self): # We need to initialize our mixin classes super(BaseSystemCollection, self).__init__() # List of system types/bases that are allowed - self.allowed_sys_types = (RodBase, RigidBodyBase, SurfaceBase) + self.allowed_sys_types:tuple[Type[SystemType], ...] = (RodBase, RigidBodyBase, SurfaceBase) # List of systems to be integrated self._systems = [] # Flag Finalize: Finalizing twice will cause an error, @@ -98,11 +99,11 @@ def insert(self, idx, system): def __str__(self): return str(self._systems) - def extend_allowed_types(self, additional_types): - self.allowed_sys_types += additional_types + def extend_allowed_types(self, additional_types: list[Type[SystemType], ...]): + self.allowed_sys_types = tuple(list(self.allowed_sys_types)+additional_types) - def override_allowed_types(self, allowed_types): - self.allowed_sys_types = allowed_types + def override_allowed_types(self, allowed_types: list[Type[SystemType], ...]): + self.allowed_sys_types = tuple(allowed_types) def _get_sys_idx_if_valid(self, sys_to_be_added): from numpy import int_ as npint From 7dcf23895f52a1b74fc18e1124fb447d3a95e6ca Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Thu, 2 May 2024 15:01:37 -0500 Subject: [PATCH 029/134] feat: allow single system integration --- elastica/timestepper/__init__.py | 29 ++++++++++++--------- elastica/timestepper/explicit_steppers.py | 13 +++++++++ elastica/timestepper/symplectic_steppers.py | 20 ++++++++++++++ 3 files changed, 50 insertions(+), 12 deletions(-) diff --git a/elastica/timestepper/__init__.py b/elastica/timestepper/__init__.py index d385edc28..2e60afa7d 100644 --- a/elastica/timestepper/__init__.py +++ b/elastica/timestepper/__init__.py @@ -11,14 +11,14 @@ from .symplectic_steppers import PositionVerlet, PEFRL from .explicit_steppers import RungeKutta4, EulerForward -from .tag import SymplecticStepperTag, ExplicitStepperTag +from .tag import SymplecticStepperTag, ExplicitStepperTag, allowed_stepper_tags from .protocol import StepperProtocol, SymplecticStepperProtocol # Deprecated: Remove in the future version # Many script still uses this method to control timestep. Keep it for backward compatibility def extend_stepper_interface( - Stepper: StepperProtocol, System: SystemType | SystemCollectionType + stepper: StepperProtocol, system_collection: SystemCollectionType ) -> Tuple[ Callable[ [StepperProtocol, SystemCollectionType, np.floating, np.floating], np.floating @@ -26,16 +26,18 @@ def extend_stepper_interface( SteppersOperatorsType, ]: # Check if system is a "collection" of smaller systems - assert is_system_a_collection(System) + assert is_system_a_collection(system_collection), "Only system-collection type can be used for timestepping. Use BaseSystemCollection." + if not hasattr(stepper, "Tag") or stepper.Tag not in allowed_stepper_tags: + raise NotImplementedError(f"{stepper} steppers is not supported. Only {allowed_stepper_tags} steppers are supported") - stepper_methods: SteppersOperatorsType = Stepper.step_methods() - do_step_method: Callable = Stepper.do_step # type: ignore[attr-defined] + stepper_methods: SteppersOperatorsType = stepper.steps_and_prefactors + do_step_method: Callable = stepper.do_step # type: ignore[attr-defined] return do_step_method, stepper_methods def integrate( - Stepper: StepperProtocol, - SystemCollection: SystemCollectionType, + stepper: StepperProtocol, + systems: SystemType | SystemCollectionType, final_time: float, n_steps: int = 1000, restart_time: float = 0.0, @@ -45,9 +47,9 @@ def integrate( Parameters ---------- - Stepper : StepperProtocol + stepper : StepperProtocol Stepper algorithm to use. - SystemCollection : SystemType + systems : SystemType | SystemCollectionType The elastica-system to simulate. final_time : float Total simulation time. The timestep is determined by final_time / n_steps. @@ -60,13 +62,16 @@ def integrate( """ assert final_time > 0.0, "Final time is negative!" assert n_steps > 0, "Number of integration steps is negative!" - assert is_system_a_collection(SystemCollection) dt = np.float_(float(final_time) / n_steps) time = np.float_(restart_time) - for i in tqdm(range(n_steps), disable=(not progress_bar)): - time = Stepper.step(SystemCollection, time, dt) + if is_system_a_collection(systems): + for i in tqdm(range(n_steps), disable=(not progress_bar)): + time = stepper.step(systems, time, dt) + else: + for i in tqdm(range(n_steps), disable=(not progress_bar)): + time = stepper.step_single_instance(systems, time, dt) print("Final time of simulation is : ", time) return time diff --git a/elastica/timestepper/explicit_steppers.py b/elastica/timestepper/explicit_steppers.py index 3064f00e1..5a336c9b1 100644 --- a/elastica/timestepper/explicit_steppers.py +++ b/elastica/timestepper/explicit_steppers.py @@ -6,6 +6,7 @@ from copy import copy from elastica.typing import ( + SystemType, SystemCollectionType, OperatorType, SteppersOperatorsType, @@ -150,6 +151,18 @@ def do_step( ) return time + def step_single_instance( + self: ExplicitStepperProtocol, + System: SystemType, + Memory: MemoryProtocol, + time: np.floating, + dt: np.floating, + ) -> np.floating: + for stage, update in self.steps_and_prefactors: + stage(System, Memory, time, dt) + time = update(System, Memory, time, dt) + return time + class EulerForward(ExplicitStepperMixin): """ diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index a09124594..c223a911b 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -5,6 +5,7 @@ from itertools import zip_longest from elastica.typing import ( + SystemType, SystemCollectionType, # StepOperatorType, # PrefactorOperatorType, @@ -138,6 +139,25 @@ def do_step( return time + def step_single_instance( + self: SymplecticStepperProtocol, + System: SystemType, + time: np.floating, + dt: np.floating, + ) -> np.floating: + + for (kin_prefactor, kin_step, dyn_step) in self.steps_and_prefactors[:-1]: + kin_step(System, time, dt) + time += kin_prefactor(dt) + System.update_internal_forces_and_torques(time) + dyn_step(System, time, dt) + + # Peel the last kinematic step and prefactor alone + last_kin_prefactor = self.steps_and_prefactors[-1][0] + last_kin_step = self.steps_and_prefactors[-1][1] + + last_kin_step(System, time, dt) + return time + last_kin_prefactor(dt) # type: ignore[no-any-return] class PositionVerlet(SymplecticStepperMixin): """ From b615cd42e3cd51722d90294440bfcd9551a98bf7 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Thu, 2 May 2024 15:14:03 -0500 Subject: [PATCH 030/134] test: fix broken tests due to changes in steppig method --- elastica/modules/base_system.py | 8 +- elastica/systems/analytical.py | 4 +- elastica/systems/protocol.py | 4 +- elastica/timestepper/__init__.py | 18 +- elastica/timestepper/explicit_steppers.py | 16 +- elastica/timestepper/protocol.py | 8 +- elastica/timestepper/symplectic_steppers.py | 33 +-- elastica/timestepper/tag.py | 2 + tests/test_math/test_timestepper.py | 251 +++++++++----------- 9 files changed, 164 insertions(+), 180 deletions(-) diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index 4c3d46528..8cbc918e0 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -58,7 +58,7 @@ def __init__(self): # We need to initialize our mixin classes super(BaseSystemCollection, self).__init__() # List of system types/bases that are allowed - self.allowed_sys_types:tuple[Type[SystemType], ...] = (RodBase, RigidBodyBase, SurfaceBase) + self.allowed_sys_types: tuple[Type, ...] = (RodBase, RigidBodyBase, SurfaceBase) # List of systems to be integrated self._systems = [] # Flag Finalize: Finalizing twice will cause an error, @@ -99,10 +99,10 @@ def insert(self, idx, system): def __str__(self): return str(self._systems) - def extend_allowed_types(self, additional_types: list[Type[SystemType], ...]): - self.allowed_sys_types = tuple(list(self.allowed_sys_types)+additional_types) + def extend_allowed_types(self, additional_types: list[Type, ...]): + self.allowed_sys_types += additional_types - def override_allowed_types(self, allowed_types: list[Type[SystemType], ...]): + def override_allowed_types(self, allowed_types: list[Type, ...]): self.allowed_sys_types = tuple(allowed_types) def _get_sys_idx_if_valid(self, sys_to_be_added): diff --git a/elastica/systems/analytical.py b/elastica/systems/analytical.py index 1b64a570d..b03a917aa 100644 --- a/elastica/systems/analytical.py +++ b/elastica/systems/analytical.py @@ -3,6 +3,7 @@ import numpy as np from elastica._rotations import _rotate from elastica.rod.data_structures import _RodSymplecticStepperMixin +from elastica.rod.rod_base import RodBase class BaseStatefulSystem: @@ -355,8 +356,9 @@ def make_simple_system_with_positions_directors( ) -class SimpleSystemWithPositionsDirectors(_RodSymplecticStepperMixin): +class SimpleSystemWithPositionsDirectors(_RodSymplecticStepperMixin, RodBase): def __init__(self, start_position, end_position, start_director): + self.ring_rod_flag = False # TODO: self.a = 0.5 self.b = 1 self.c = 2 diff --git a/elastica/systems/protocol.py b/elastica/systems/protocol.py index c2864f608..3d78f32e0 100644 --- a/elastica/systems/protocol.py +++ b/elastica/systems/protocol.py @@ -38,6 +38,8 @@ def external_forces(self) -> NDArray: ... @property def external_torques(self) -> NDArray: ... + def update_internal_forces_and_torques(self, time: np.floating) -> None: ... + class SymplecticSystemProtocol(SystemProtocol, Protocol): """ @@ -64,8 +66,6 @@ def dynamic_rates( self, time: np.floating, prefac: np.floating ) -> tuple[NDArray]: ... - def update_internal_forces_and_torques(self, time: np.floating) -> None: ... - class ExplicitSystemProtocol(SystemProtocol, Protocol): # TODO: Temporarily made to handle explicit stepper. diff --git a/elastica/timestepper/__init__.py b/elastica/timestepper/__init__.py index 2e60afa7d..49876dd43 100644 --- a/elastica/timestepper/__init__.py +++ b/elastica/timestepper/__init__.py @@ -1,6 +1,6 @@ __doc__ = """Timestepping utilities to be used with Rod and RigidBody classes""" -from typing import Tuple, List, Callable, Type, Any +from typing import Tuple, List, Callable, Type, Any, overload from elastica.typing import SystemType, SystemCollectionType, SteppersOperatorsType import numpy as np @@ -26,9 +26,13 @@ def extend_stepper_interface( SteppersOperatorsType, ]: # Check if system is a "collection" of smaller systems - assert is_system_a_collection(system_collection), "Only system-collection type can be used for timestepping. Use BaseSystemCollection." + assert is_system_a_collection( + system_collection + ), "Only system-collection type can be used for timestepping. Use BaseSystemCollection." if not hasattr(stepper, "Tag") or stepper.Tag not in allowed_stepper_tags: - raise NotImplementedError(f"{stepper} steppers is not supported. Only {allowed_stepper_tags} steppers are supported") + raise NotImplementedError( + f"{stepper} steppers is not supported. Only {allowed_stepper_tags} steppers are supported" + ) stepper_methods: SteppersOperatorsType = stepper.steps_and_prefactors do_step_method: Callable = stepper.do_step # type: ignore[attr-defined] @@ -42,7 +46,7 @@ def integrate( n_steps: int = 1000, restart_time: float = 0.0, progress_bar: bool = True, -) -> np.floating: +) -> float: """ Parameters @@ -68,10 +72,10 @@ def integrate( if is_system_a_collection(systems): for i in tqdm(range(n_steps), disable=(not progress_bar)): - time = stepper.step(systems, time, dt) + time = stepper.step(systems, time, dt) # type: ignore[arg-type] else: for i in tqdm(range(n_steps), disable=(not progress_bar)): - time = stepper.step_single_instance(systems, time, dt) + time = stepper.step_single_instance(systems, time, dt) # type: ignore[arg-type] print("Final time of simulation is : ", time) - return time + return float(time) diff --git a/elastica/timestepper/explicit_steppers.py b/elastica/timestepper/explicit_steppers.py index 5a336c9b1..00afbaf6b 100644 --- a/elastica/timestepper/explicit_steppers.py +++ b/elastica/timestepper/explicit_steppers.py @@ -133,22 +133,20 @@ def step( @staticmethod def do_step( TimeStepper: ExplicitStepperProtocol, - stages_and_updates: SteppersOperatorsType, + steps_and_prefactors: SteppersOperatorsType, SystemCollection: SystemCollectionType, MemoryCollection: Any, # TODO time: np.floating, dt: np.floating, ) -> np.floating: - for stage, update in stages_and_updates: + for stage, update in steps_and_prefactors: SystemCollection.synchronize(time) for system, memory in zip(SystemCollection[:-1], MemoryCollection[:-1]): - stage(TimeStepper, system, memory, time, dt) - _ = update(TimeStepper, system, memory, time, dt) + stage(system, memory, time, dt) + _ = update(system, memory, time, dt) - stage(TimeStepper, SystemCollection[-1], MemoryCollection[-1], time, dt) - time = update( - TimeStepper, SystemCollection[-1], MemoryCollection[-1], time, dt - ) + stage(SystemCollection[-1], MemoryCollection[-1], time, dt) + time = update(SystemCollection[-1], MemoryCollection[-1], time, dt) return time def step_single_instance( @@ -197,7 +195,7 @@ def _first_update( return time + dt -class RungeKutta4: +class RungeKutta4(ExplicitStepperMixin): """ Stateless runge-kutta4. coordinates operations only, memory needs to be externally managed and allocated. diff --git a/elastica/timestepper/protocol.py b/elastica/timestepper/protocol.py index f923e62eb..d27279269 100644 --- a/elastica/timestepper/protocol.py +++ b/elastica/timestepper/protocol.py @@ -17,21 +17,23 @@ class StepperProtocol(Protocol): """Protocol for all time-steppers""" Tag: StepperTags + steps_and_prefactors: SteppersOperatorsType def __init__(self) -> None: ... @property def n_stages(self) -> int: ... - @property - def steps_and_prefactors(self) -> SteppersOperatorsType: ... - def step_methods(self) -> SteppersOperatorsType: ... def step( self, SystemCollection: SystemCollectionType, time: np.floating, dt: np.floating ) -> np.floating: ... + def step_single_instance( + self, SystemCollection: SystemType, time: np.floating, dt: np.floating + ) -> np.floating: ... + class SymplecticStepperProtocol(StepperProtocol, Protocol): """symplectic stepper protocol.""" diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index c223a911b..8e38e279a 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -34,21 +34,18 @@ class SymplecticStepperMixin: def __init__(self: SymplecticStepperProtocol): - self.steps_and_prefactors: Final = self.step_methods() + self.steps_and_prefactors: Final[SteppersOperatorsType] = self.step_methods() def step_methods(self: SymplecticStepperProtocol) -> SteppersOperatorsType: # Let the total number of steps for the Symplectic method - # be (2*n + 1) (for time-symmetry). What we do is collect - # the first n + 1 entries down in _steps and _prefac below, and then - # reverse and append it to itself. + # be (2*n + 1) (for time-symmetry). _steps: list[OperatorType] = self.get_steps() - _steps = _steps + _steps[-2::-1] - # Prefac here is necessary because the linear-exponential integrator # needs only the prefactor and not the dt. _prefactors: list[OperatorType] = self.get_prefactors() - _prefactors = _prefactors + _prefactors[-2::-1] - assert len(_steps) == 2 * len(_prefactors) - 1 + assert int(np.ceil(len(_steps) / 2)) == len( + _prefactors + ), f"{len(_steps)=}, {len(_prefactors)=}" # Separate the kinematic and dynamic steps _kinematic_steps: list[OperatorType] = _steps[::2] @@ -100,9 +97,9 @@ def do_step( for kin_prefactor, kin_step, dyn_step in steps_and_prefactors[:-1]: for system in SystemCollection._memory_blocks: - kin_step(TimeStepper, system, time, dt) + kin_step(system, time, dt) - time += kin_prefactor(TimeStepper, dt) + time += kin_prefactor(dt) # Constrain only values SystemCollection.constrain_values(time) @@ -116,7 +113,7 @@ def do_step( SystemCollection.synchronize(time) for system in SystemCollection._memory_blocks: - dyn_step(TimeStepper, system, time, dt) + dyn_step(system, time, dt) # Constrain only rates SystemCollection.constrain_rates(time) @@ -126,8 +123,8 @@ def do_step( last_kin_step = steps_and_prefactors[-1][1] for system in SystemCollection._memory_blocks: - last_kin_step(TimeStepper, system, time, dt) - time += last_kin_prefactor(TimeStepper, dt) + last_kin_step(system, time, dt) + time += last_kin_prefactor(dt) SystemCollection.constrain_values(time) # Call back function, will call the user defined call back functions and store data @@ -146,7 +143,7 @@ def step_single_instance( dt: np.floating, ) -> np.floating: - for (kin_prefactor, kin_step, dyn_step) in self.steps_and_prefactors[:-1]: + for kin_prefactor, kin_step, dyn_step in self.steps_and_prefactors[:-1]: kin_step(System, time, dt) time += kin_prefactor(dt) System.update_internal_forces_and_torques(time) @@ -159,6 +156,7 @@ def step_single_instance( last_kin_step(System, time, dt) return time + last_kin_prefactor(dt) # type: ignore[no-any-return] + class PositionVerlet(SymplecticStepperMixin): """ Position Verlet symplectic time stepper class, which @@ -171,11 +169,13 @@ def get_steps(self) -> list[OperatorType]: return [ self._first_kinematic_step, self._first_dynamic_step, + self._first_kinematic_step, ] def get_prefactors(self) -> list[OperatorType]: return [ self._first_prefactor, + self._first_prefactor, ] def _first_prefactor(self, dt: np.floating) -> np.floating: @@ -222,19 +222,22 @@ class PEFRL(SymplecticStepperMixin): xi_chi_dash_coeff: np.float64 = 1.0 - 2.0 * (ξ + χ) def get_steps(self) -> list[OperatorType]: - return [ + operators = [ self._first_kinematic_step, self._first_dynamic_step, self._second_kinematic_step, self._second_dynamic_step, self._third_kinematic_step, ] + return operators + operators[-2::-1] def get_prefactors(self) -> list[OperatorType]: return [ self._first_kinematic_prefactor, self._second_kinematic_prefactor, self._third_kinematic_prefactor, + self._second_kinematic_prefactor, + self._first_kinematic_prefactor, ] def _first_kinematic_prefactor(self, dt: np.floating) -> np.floating: diff --git a/elastica/timestepper/tag.py b/elastica/timestepper/tag.py index d23788750..8f0d77517 100644 --- a/elastica/timestepper/tag.py +++ b/elastica/timestepper/tag.py @@ -4,3 +4,5 @@ ExplicitStepperTag: Final = "ExplicitStepper" SymplecticStepperTag: Final = "SymplecticStepper" StepperTags: TypeAlias = Literal["SymplecticStepper", "ExplicitStepper"] + +allowed_stepper_tags: list[StepperTags] = [ExplicitStepperTag, SymplecticStepperTag] diff --git a/tests/test_math/test_timestepper.py b/tests/test_math/test_timestepper.py index 3868013f9..69c533e9c 100644 --- a/tests/test_math/test_timestepper.py +++ b/tests/test_math/test_timestepper.py @@ -37,22 +37,34 @@ class TestExtendStepperInterface: class MockSymplecticStepper(SymplecticStepperMixin): Tag = SymplecticStepperTag - def _first_prefactor(self): + def get_steps(self): + return [self._kinematic_step, self._dynamic_step, self._kinematic_step] + + def get_prefactors(self): + return [self._prefactor, self._prefactor] + + def _prefactor(self): pass - def _first_kinematic_step(self): + def _kinematic_step(self): pass - def _first_dynamic_step(self): + def _dynamic_step(self): pass class MockExplicitStepper(ExplicitStepperMixin): Tag = ExplicitStepperTag - def _first_stage(self): + def get_stages(self): + return [self._stage] + + def get_updates(self): + return [self._update] + + def _stage(self): pass - def _first_update(self): + def _update(self): pass # We cannot call a stepper on a system until both the stepper @@ -64,28 +76,33 @@ def _first_update(self): "stepper_module", [MockSymplecticStepper, MockExplicitStepper], ) - def test_symplectic_stepper_interface_for_collective_systems( - self, stepper_module - ): + def test_symplectic_stepper_interface_for_collective_systems(self, stepper_module): system = SymplecticUndampedHarmonicOscillatorCollectiveSystem() - stepper_cls = stepper_module.steps_and_prefactors - stepper = stepper_cls() + stepper = stepper_module() stepper_methods = None - assert stepper_methods is None - _, stepper_methods = extend_stepper_interface(stepper, system) - assert stepper_methods + assert stepper_methods == stepper.steps_and_prefactors class MockBadStepper: Tag = int() # an arbitrary tag that doesn't mean anything - @pytest.mark.parametrize("stepper_and_interface", [MockBadStepper]) - def test_symplectic_stepper_throws_for_bad_stepper(self, stepper_and_interface): + @pytest.mark.parametrize("stepper_module", [MockSymplecticStepper]) + def test_symplectic_stepper_throws_for_bad_stepper_for_simple_system( + self, stepper_module + ): system = ScalarExponentialDecaySystem() - stepper_cls = stepper_and_interface - stepper = stepper_cls() + stepper = stepper_module() + + with pytest.raises(AssertionError) as excinfo: + extend_stepper_interface(stepper, system) + assert "Only system-collection type can be used" in str(excinfo.value) + + @pytest.mark.parametrize("stepper_module", [MockBadStepper]) + def test_symplectic_stepper_throws_for_bad_stepper(self, stepper_module): + system = SymplecticUndampedHarmonicOscillatorCollectiveSystem() + stepper = stepper_module() with pytest.raises(NotImplementedError) as excinfo: extend_stepper_interface(stepper, system) @@ -112,127 +129,79 @@ def test_integrate_throws_an_assert_for_negative_total_steps(): SymplecticSteppers = [PositionVerlet, PEFRL] -""" -class TestExplicitSteppers: - @pytest.mark.parametrize("stepper", StatefulExplicitSteppers) - def test_against_scalar_exponential(self, stepper): - system = ScalarExponentialDecaySystem(-1, 1) - final_time = 1 - n_steps = 1000 - integrate(stepper(), system, final_time=final_time, n_steps=n_steps) - - assert_allclose( - system.state, - system.analytical_solution(final_time), - rtol=Tolerance.rtol() * 1e3, - atol=Tolerance.atol(), - ) - - @pytest.mark.parametrize("stepper", StatefulExplicitSteppers[:-1]) - def test_against_undamped_harmonic_oscillator(self, stepper): - system = UndampedSimpleHarmonicOscillatorSystem() - final_time = 4.0 * np.pi - n_steps = 2000 - integrate(stepper(), system, final_time=final_time, n_steps=n_steps) - - assert_allclose( - system.state, - system.analytical_solution(final_time), - rtol=Tolerance.rtol(), - atol=Tolerance.atol(), - ) - - @pytest.mark.parametrize("stepper", StatefulExplicitSteppers[:-1]) - def test_against_damped_harmonic_oscillator(self, stepper): - system = DampedSimpleHarmonicOscillatorSystem() - final_time = 4.0 * np.pi - n_steps = 2000 - integrate(stepper(), system, final_time=final_time, n_steps=n_steps) - - assert_allclose( - system.state, - system.analytical_solution(final_time), - rtol=Tolerance.rtol(), - atol=Tolerance.atol(), - ) - - def test_linear_exponential_integrator(self): - system = MultipleFrameRotationSystem(n_frames=128) - final_time = np.pi - n_steps = 1000 - integrate( - StatefulLinearExponentialIntegrator(), - system, - final_time=final_time, - n_steps=n_steps, - ) - - assert_allclose( - system.linearly_evolving_state, - system.analytical_solution(final_time), - atol=1e-4, - ) - - @pytest.mark.parametrize("explicit_stepper", StatefulExplicitSteppers[:-1]) - def test_explicit_against_analytical_system(self, explicit_stepper): - system = SecondOrderHybridSystem() - final_time = 1.0 - n_steps = 2000 - integrate(explicit_stepper(), system, final_time=final_time, n_steps=n_steps) - - assert_allclose( - system.final_solution(final_time), - system.analytical_solution(final_time), - rtol=Tolerance.rtol() * 1e2, - atol=Tolerance.atol(), - ) -""" - - -class TestSymplecticSteppers: - @pytest.mark.parametrize("stepper", SymplecticSteppers) - def test_symplectic_against_undamped_harmonic_oscillator(self, stepper): - system = SymplecticUndampedSimpleHarmonicOscillatorSystem( - omega=1.0 * np.pi, init_val=np.array([0.2, 0.8]) - ) - final_time = 4.0 * np.pi - n_steps = 2000 - time_stepper = stepper() - integrate(time_stepper, system, final_time=final_time, n_steps=n_steps) - - # Symplectic systems conserve energy to a certain extent - assert_allclose( - *system.compute_energy(final_time), - rtol=Tolerance.rtol() * 1e1, - atol=Tolerance.atol(), - ) - - # assert_allclose( - # system._state, - # system.analytical_solution(final_time), - # rtol=Tolerance.rtol(), - # atol=Tolerance.atol(), - # ) - - -""" - @pytest.mark.xfail - @pytest.mark.parametrize("symplectic_stepper", SymplecticSteppers) - def test_hybrid_symplectic_against_analytical_system(self, symplectic_stepper): - system = SecondOrderHybridSystem() - final_time = 1.0 - n_steps = 2000 - # stepper = SymplecticCosseratRodStepper(symplectic_stepper=symplectic_stepper()) - stepper = symplectic_stepper() - integrate(stepper, system, final_time=final_time, n_steps=n_steps) - - assert_allclose( - system.final_solution(final_time), - system.analytical_solution(final_time), - rtol=Tolerance.rtol() * 1e2, - atol=Tolerance.atol(), - ) -""" +# class TestExplicitSteppers: +# @pytest.mark.parametrize("stepper", StatefulExplicitSteppers) +# def test_against_scalar_exponential(self, stepper): +# system = ScalarExponentialDecaySystem(-1, 1) +# final_time = 1 +# n_steps = 1000 +# integrate(stepper(), system, final_time=final_time, n_steps=n_steps) +# +# assert_allclose( +# system.state, +# system.analytical_solution(final_time), +# rtol=Tolerance.rtol() * 1e3, +# atol=Tolerance.atol(), +# ) +# +# @pytest.mark.parametrize("stepper", StatefulExplicitSteppers[:-1]) +# def test_against_undamped_harmonic_oscillator(self, stepper): +# system = UndampedSimpleHarmonicOscillatorSystem() +# final_time = 4.0 * np.pi +# n_steps = 2000 +# integrate(stepper(), system, final_time=final_time, n_steps=n_steps) +# +# assert_allclose( +# system.state, +# system.analytical_solution(final_time), +# rtol=Tolerance.rtol(), +# atol=Tolerance.atol(), +# ) +# +# @pytest.mark.parametrize("stepper", StatefulExplicitSteppers[:-1]) +# def test_against_damped_harmonic_oscillator(self, stepper): +# system = DampedSimpleHarmonicOscillatorSystem() +# final_time = 4.0 * np.pi +# n_steps = 2000 +# integrate(stepper(), system, final_time=final_time, n_steps=n_steps) +# +# assert_allclose( +# system.state, +# system.analytical_solution(final_time), +# rtol=Tolerance.rtol(), +# atol=Tolerance.atol(), +# ) +# +# def test_linear_exponential_integrator(self): +# system = MultipleFrameRotationSystem(n_frames=128) +# final_time = np.pi +# n_steps = 1000 +# integrate( +# StatefulLinearExponentialIntegrator(), +# system, +# final_time=final_time, +# n_steps=n_steps, +# ) +# +# assert_allclose( +# system.linearly_evolving_state, +# system.analytical_solution(final_time), +# atol=1e-4, +# ) +# +# @pytest.mark.parametrize("explicit_stepper", StatefulExplicitSteppers[:-1]) +# def test_explicit_against_analytical_system(self, explicit_stepper): +# system = SecondOrderHybridSystem() +# final_time = 1.0 +# n_steps = 2000 +# integrate(explicit_stepper(), system, final_time=final_time, n_steps=n_steps) +# +# assert_allclose( +# system.final_solution(final_time), +# system.analytical_solution(final_time), +# rtol=Tolerance.rtol() * 1e2, +# atol=Tolerance.atol(), +# ) class TestSteppersAgainstCollectiveSystems: @@ -350,9 +319,13 @@ def test_symplectics_against_ellipse_motion(self, symplectic_stepper): ) final_time = 1.0 n_steps = 1000 + dt = final_time / n_steps + stepper = symplectic_stepper() - integrate(stepper, rod_like_system, final_time=final_time, n_steps=n_steps) + time = 0.0 + for _ in range(n_steps): + time = stepper.step_single_instance(rod_like_system, time, dt) assert_allclose( rod_like_system.position_collection, From 0531ef87aec8a066e30783d5580f981faf56eaf2 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Fri, 3 May 2024 17:54:27 -0500 Subject: [PATCH 031/134] remove: timestepper tag --- elastica/timestepper/__init__.py | 12 +++++------- elastica/timestepper/explicit_steppers.py | 5 ----- elastica/timestepper/protocol.py | 2 -- elastica/timestepper/symplectic_steppers.py | 5 ----- elastica/timestepper/tag.py | 8 -------- tests/test_math/test_timestepper.py | 7 ++----- 6 files changed, 7 insertions(+), 32 deletions(-) delete mode 100644 elastica/timestepper/tag.py diff --git a/elastica/timestepper/__init__.py b/elastica/timestepper/__init__.py index 49876dd43..5b6111d09 100644 --- a/elastica/timestepper/__init__.py +++ b/elastica/timestepper/__init__.py @@ -11,7 +11,6 @@ from .symplectic_steppers import PositionVerlet, PEFRL from .explicit_steppers import RungeKutta4, EulerForward -from .tag import SymplecticStepperTag, ExplicitStepperTag, allowed_stepper_tags from .protocol import StepperProtocol, SymplecticStepperProtocol @@ -29,13 +28,12 @@ def extend_stepper_interface( assert is_system_a_collection( system_collection ), "Only system-collection type can be used for timestepping. Use BaseSystemCollection." - if not hasattr(stepper, "Tag") or stepper.Tag not in allowed_stepper_tags: - raise NotImplementedError( - f"{stepper} steppers is not supported. Only {allowed_stepper_tags} steppers are supported" - ) - stepper_methods: SteppersOperatorsType = stepper.steps_and_prefactors - do_step_method: Callable = stepper.do_step # type: ignore[attr-defined] + try: + stepper_methods: SteppersOperatorsType = stepper.steps_and_prefactors + do_step_method: Callable = stepper.do_step # type: ignore[attr-defined] + except AttributeError as e: + raise NotImplementedError(f"{stepper} stepper is not supported.") from e return do_step_method, stepper_methods diff --git a/elastica/timestepper/explicit_steppers.py b/elastica/timestepper/explicit_steppers.py index 00afbaf6b..56140e111 100644 --- a/elastica/timestepper/explicit_steppers.py +++ b/elastica/timestepper/explicit_steppers.py @@ -13,7 +13,6 @@ StateType, ) from elastica.systems.protocol import ExplicitSystemProtocol -from .tag import StepperTags, ExplicitStepperTag from .protocol import ExplicitStepperProtocol, MemoryProtocol @@ -167,8 +166,6 @@ class EulerForward(ExplicitStepperMixin): Classical Euler Forward stepper. Stateless, coordinates operations only. """ - Tag: StepperTags = ExplicitStepperTag - def get_stages(self) -> list[OperatorType]: return [self._first_stage] @@ -201,8 +198,6 @@ class RungeKutta4(ExplicitStepperMixin): to be externally managed and allocated. """ - Tag: StepperTags = ExplicitStepperTag - def get_stages(self) -> list[OperatorType]: return [ self._first_stage, diff --git a/elastica/timestepper/protocol.py b/elastica/timestepper/protocol.py index d27279269..818d04b8b 100644 --- a/elastica/timestepper/protocol.py +++ b/elastica/timestepper/protocol.py @@ -8,7 +8,6 @@ OperatorType, SystemCollectionType, ) -from .tag import StepperTags import numpy as np @@ -16,7 +15,6 @@ class StepperProtocol(Protocol): """Protocol for all time-steppers""" - Tag: StepperTags steps_and_prefactors: SteppersOperatorsType def __init__(self) -> None: ... diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index 8e38e279a..de2d761f1 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -21,7 +21,6 @@ ) from elastica.systems.protocol import SymplecticSystemProtocol from .protocol import SymplecticStepperProtocol -from .tag import StepperTags, SymplecticStepperTag """ Developer Note @@ -163,8 +162,6 @@ class PositionVerlet(SymplecticStepperMixin): includes methods for second-order position Verlet. """ - Tag: StepperTags = SymplecticStepperTag - def get_steps(self) -> list[OperatorType]: return [ self._first_kinematic_step, @@ -210,8 +207,6 @@ class PEFRL(SymplecticStepperMixin): http://arxiv.org/abs/cond-mat/0110585 """ - Tag: StepperTags = SymplecticStepperTag - # xi and chi are confusing, but be careful! ξ: np.float64 = np.float64(0.1786178958448091e0) # ξ λ: np.float64 = -np.float64(0.2123418310626054e0) # λ diff --git a/elastica/timestepper/tag.py b/elastica/timestepper/tag.py deleted file mode 100644 index 8f0d77517..000000000 --- a/elastica/timestepper/tag.py +++ /dev/null @@ -1,8 +0,0 @@ -from typing import Literal, TypeAlias, Final - - -ExplicitStepperTag: Final = "ExplicitStepper" -SymplecticStepperTag: Final = "SymplecticStepper" -StepperTags: TypeAlias = Literal["SymplecticStepper", "ExplicitStepper"] - -allowed_stepper_tags: list[StepperTags] = [ExplicitStepperTag, SymplecticStepperTag] diff --git a/tests/test_math/test_timestepper.py b/tests/test_math/test_timestepper.py index 69c533e9c..11dbfa8d3 100644 --- a/tests/test_math/test_timestepper.py +++ b/tests/test_math/test_timestepper.py @@ -25,7 +25,6 @@ PEFRL, SymplecticStepperMixin, ) -from elastica.timestepper.tag import SymplecticStepperTag, ExplicitStepperTag from elastica.utils import Tolerance @@ -35,7 +34,6 @@ class TestExtendStepperInterface: """TODO add documentation""" class MockSymplecticStepper(SymplecticStepperMixin): - Tag = SymplecticStepperTag def get_steps(self): return [self._kinematic_step, self._dynamic_step, self._kinematic_step] @@ -53,7 +51,6 @@ def _dynamic_step(self): pass class MockExplicitStepper(ExplicitStepperMixin): - Tag = ExplicitStepperTag def get_stages(self): return [self._stage] @@ -86,7 +83,7 @@ def test_symplectic_stepper_interface_for_collective_systems(self, stepper_modul assert stepper_methods == stepper.steps_and_prefactors class MockBadStepper: - Tag = int() # an arbitrary tag that doesn't mean anything + pass @pytest.mark.parametrize("stepper_module", [MockSymplecticStepper]) def test_symplectic_stepper_throws_for_bad_stepper_for_simple_system( @@ -106,7 +103,7 @@ def test_symplectic_stepper_throws_for_bad_stepper(self, stepper_module): with pytest.raises(NotImplementedError) as excinfo: extend_stepper_interface(stepper, system) - assert "steppers are supported" in str(excinfo.value) + assert "stepper is not supported" in str(excinfo.value) def test_integrate_throws_an_assert_for_negative_final_time(): From 425e14a023a8bc393cf0f4b7c418b22a050dc212 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Fri, 3 May 2024 17:57:11 -0500 Subject: [PATCH 032/134] update tutorial with new timestepper syntax --- examples/Binder/1_Timoshenko_Beam.ipynb | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/examples/Binder/1_Timoshenko_Beam.ipynb b/examples/Binder/1_Timoshenko_Beam.ipynb index a8b81fdfa..63efe00db 100644 --- a/examples/Binder/1_Timoshenko_Beam.ipynb +++ b/examples/Binder/1_Timoshenko_Beam.ipynb @@ -627,16 +627,14 @@ "\n", "\n", "def run_and_update_plot(simulator, dt, start_time, stop_time, ax):\n", - " from elastica.timestepper import extend_stepper_interface\n", " from elastica.timestepper.symplectic_steppers import PositionVerlet\n", "\n", " timestepper = PositionVerlet()\n", - " do_step, stages_and_updates = extend_stepper_interface(timestepper, simulator)\n", "\n", " n_steps = int((stop_time - start_time) / dt)\n", " time = start_time\n", " for i in range(n_steps):\n", - " time = do_step(timestepper, stages_and_updates, simulator, time, dt)\n", + " time = timestepper.step(simulator, time, dt)\n", " plot_timoshenko_dynamic(shearable_rod_new, unshearable_rod_new, end_force, time, ax)\n", " return time\n", "\n", From 52a77da035088b689fbaadeacd317cc459e26473 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Fri, 3 May 2024 23:09:32 -0500 Subject: [PATCH 033/134] feat: add testing for single system integrate overload --- elastica/timestepper/__init__.py | 28 ++++++++++++++++++++++------ tests/test_math/test_timestepper.py | 27 ++++++++++++++++----------- 2 files changed, 38 insertions(+), 17 deletions(-) diff --git a/elastica/timestepper/__init__.py b/elastica/timestepper/__init__.py index 5b6111d09..34575dfd6 100644 --- a/elastica/timestepper/__init__.py +++ b/elastica/timestepper/__init__.py @@ -10,7 +10,6 @@ from .symplectic_steppers import PositionVerlet, PEFRL from .explicit_steppers import RungeKutta4, EulerForward - from .protocol import StepperProtocol, SymplecticStepperProtocol @@ -24,11 +23,6 @@ def extend_stepper_interface( ], SteppersOperatorsType, ]: - # Check if system is a "collection" of smaller systems - assert is_system_a_collection( - system_collection - ), "Only system-collection type can be used for timestepping. Use BaseSystemCollection." - try: stepper_methods: SteppersOperatorsType = stepper.steps_and_prefactors do_step_method: Callable = stepper.do_step # type: ignore[attr-defined] @@ -37,6 +31,28 @@ def extend_stepper_interface( return do_step_method, stepper_methods +@overload +def integrate( + stepper: StepperProtocol, + systems: SystemType, + final_time: float, + n_steps: int, + restart_time: float, + progress_bar: bool, +) -> float: ... + + +@overload +def integrate( + stepper: StepperProtocol, + systems: SystemCollectionType, + final_time: float, + n_steps: int, + restart_time: float, + progress_bar: bool, +) -> float: ... + + def integrate( stepper: StepperProtocol, systems: SystemType | SystemCollectionType, diff --git a/tests/test_math/test_timestepper.py b/tests/test_math/test_timestepper.py index 11dbfa8d3..bf4b2fa2c 100644 --- a/tests/test_math/test_timestepper.py +++ b/tests/test_math/test_timestepper.py @@ -69,6 +69,22 @@ def _update(self): # So before "seeing" the system, the stepper should not have # the interface (interface_cls). It should however have the interface # after "seeing" the system, via extend_stepper_interface + @pytest.mark.parametrize( + "stepper_module", + [ + MockSymplecticStepper, + MockExplicitStepper, + ], + ) + def test_symplectic_stepper_interface_for_simple_systems(self, stepper_module): + system = ScalarExponentialDecaySystem() + stepper = stepper_module() + + stepper_methods = None + _, stepper_methods = extend_stepper_interface(stepper, system) + + assert stepper_methods + @pytest.mark.parametrize( "stepper_module", [MockSymplecticStepper, MockExplicitStepper], @@ -85,17 +101,6 @@ def test_symplectic_stepper_interface_for_collective_systems(self, stepper_modul class MockBadStepper: pass - @pytest.mark.parametrize("stepper_module", [MockSymplecticStepper]) - def test_symplectic_stepper_throws_for_bad_stepper_for_simple_system( - self, stepper_module - ): - system = ScalarExponentialDecaySystem() - stepper = stepper_module() - - with pytest.raises(AssertionError) as excinfo: - extend_stepper_interface(stepper, system) - assert "Only system-collection type can be used" in str(excinfo.value) - @pytest.mark.parametrize("stepper_module", [MockBadStepper]) def test_symplectic_stepper_throws_for_bad_stepper(self, stepper_module): system = SymplecticUndampedHarmonicOscillatorCollectiveSystem() From 42bfd091903dcc0ea26aa5af6ff2d7c6bb230599 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sat, 4 May 2024 16:52:47 -0500 Subject: [PATCH 034/134] autoflake8 -> autoflake --- Makefile | 24 ++++---- elastica/timestepper/explicit_steppers.py | 2 +- elastica/timestepper/protocol.py | 2 +- elastica/timestepper/symplectic_steppers.py | 2 +- poetry.lock | 66 ++++++++++++--------- pyproject.toml | 11 +++- 6 files changed, 62 insertions(+), 45 deletions(-) diff --git a/Makefile b/Makefile index abb6fd70f..89b4b1799 100644 --- a/Makefile +++ b/Makefile @@ -1,7 +1,7 @@ #* Variables PYTHON := python3 PYTHONPATH := `pwd` -AUTOFLAKE8_ARGS := -r --exclude '__init__.py' --keep-pass-after-docstring +AUTOFLAKE_ARGS := -r #* Poetry .PHONY: poetry-download poetry-download: @@ -47,19 +47,19 @@ flake8: poetry run flake8 --version poetry run flake8 elastica tests -.PHONY: autoflake8-check -autoflake8-check: - poetry run autoflake8 --version - poetry run autoflake8 $(AUTOFLAKE8_ARGS) elastica tests examples - poetry run autoflake8 --check $(AUTOFLAKE8_ARGS) elastica tests examples +.PHONY: autoflake-check +autoflake-check: + poetry run autoflake --version + poetry run autoflake $(AUTOFLAKE_ARGS) elastica tests examples + poetry run autoflake --check $(AUTOFLAKE_ARGS) elastica tests examples -.PHONY: autoflake8-format -autoflake8-format: - poetry run autoflake8 --version - poetry run autoflake8 --in-place $(AUTOFLAKE8_ARGS) elastica tests examples +.PHONY: autoflake-format +autoflake-format: + poetry run autoflake --version + poetry run autoflake --in-place $(AUTOFLAKE_ARGS) elastica tests examples .PHONY: format-codestyle -format-codestyle: black flake8 +format-codestyle: black autoflake-format .PHONY: mypy mypy: @@ -78,7 +78,7 @@ test_coverage_xml: NUMBA_DISABLE_JIT=1 poetry run pytest --cov=elastica --cov-report=xml .PHONY: check-codestyle -check-codestyle: black-check flake8 autoflake8-check +check-codestyle: black-check flake8 autoflake-check .PHONY: formatting formatting: format-codestyle diff --git a/elastica/timestepper/explicit_steppers.py b/elastica/timestepper/explicit_steppers.py index 56140e111..aeb4cab03 100644 --- a/elastica/timestepper/explicit_steppers.py +++ b/elastica/timestepper/explicit_steppers.py @@ -1,6 +1,6 @@ __doc__ = """Explicit timesteppers and concepts""" -from typing import Type, Any, Final +from typing import Any import numpy as np from copy import copy diff --git a/elastica/timestepper/protocol.py b/elastica/timestepper/protocol.py index 818d04b8b..6eff40dd9 100644 --- a/elastica/timestepper/protocol.py +++ b/elastica/timestepper/protocol.py @@ -1,6 +1,6 @@ __doc__ = "Time stepper interface" -from typing import Protocol, Callable, Literal, ClassVar, Type +from typing import Protocol from elastica.typing import ( SystemType, diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index de2d761f1..84a87ce20 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -1,6 +1,6 @@ __doc__ = """Symplectic time steppers and concepts for integrating the kinematic and dynamic equations of rod-like objects. """ -from typing import Callable, Any, Final +from typing import Any, Final from itertools import zip_longest diff --git a/poetry.lock b/poetry.lock index 7883be3f5..282f244d9 100644 --- a/poetry.lock +++ b/poetry.lock @@ -26,18 +26,19 @@ files = [ ] [[package]] -name = "autoflake8" -version = "0.4.1" -description = "Tool to automatically fix some issues reported by flake8 (forked from autoflake)." +name = "autoflake" +version = "2.3.1" +description = "Removes unused imports and unused variables" optional = false -python-versions = ">=3.7,<4.0" +python-versions = ">=3.8" files = [ - {file = "autoflake8-0.4.1-py3-none-any.whl", hash = "sha256:fdf663b627993ac38e5b55b7d742c388fb2a4f34798a052f43eecc5e8d629e9d"}, - {file = "autoflake8-0.4.1.tar.gz", hash = "sha256:c17da499bd2b71ba02fb11fe53ff1ad83d7dae6efb0f115fd1344f467797c679"}, + {file = "autoflake-2.3.1-py3-none-any.whl", hash = "sha256:3ae7495db9084b7b32818b4140e6dc4fc280b712fb414f5b8fe57b0a8e85a840"}, + {file = "autoflake-2.3.1.tar.gz", hash = "sha256:c98b75dc5b0a86459c4f01a1d32ac7eb4338ec4317a4469515ff1e687ecd909e"}, ] [package.dependencies] -pyflakes = ">=2.3.0" +pyflakes = ">=3.0.0" +tomli = {version = ">=2.0.1", markers = "python_version < \"3.11\""} [[package]] name = "babel" @@ -496,19 +497,19 @@ typing = ["typing-extensions (>=4.8)"] [[package]] name = "flake8" -version = "3.9.2" +version = "7.0.0" description = "the modular source code checker: pep8 pyflakes and co" optional = false -python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,>=2.7" +python-versions = ">=3.8.1" files = [ - {file = "flake8-3.9.2-py2.py3-none-any.whl", hash = "sha256:bf8fd333346d844f616e8d47905ef3a3384edae6b4e9beb0c5101e25e3110907"}, - {file = "flake8-3.9.2.tar.gz", hash = "sha256:07528381786f2a6237b061f6e96610a4167b226cb926e2aa2b6b1d78057c576b"}, + {file = "flake8-7.0.0-py2.py3-none-any.whl", hash = "sha256:a6dfbb75e03252917f2473ea9653f7cd799c3064e54d4c8140044c5c065f53c3"}, + {file = "flake8-7.0.0.tar.gz", hash = "sha256:33f96621059e65eec474169085dc92bf26e7b2d47366b70be2f67ab80dc25132"}, ] [package.dependencies] -mccabe = ">=0.6.0,<0.7.0" -pycodestyle = ">=2.7.0,<2.8.0" -pyflakes = ">=2.3.0,<2.4.0" +mccabe = ">=0.7.0,<0.8.0" +pycodestyle = ">=2.11.0,<2.12.0" +pyflakes = ">=3.2.0,<3.3.0" [[package]] name = "fonttools" @@ -928,13 +929,13 @@ python-dateutil = ">=2.7" [[package]] name = "mccabe" -version = "0.6.1" +version = "0.7.0" description = "McCabe checker, plugin for flake8" optional = false -python-versions = "*" +python-versions = ">=3.6" files = [ - {file = "mccabe-0.6.1-py2.py3-none-any.whl", hash = "sha256:ab8a6258860da4b6677da4bd2fe5dc2c659cff31b3ee4f7f5d64e79735b80d42"}, - {file = "mccabe-0.6.1.tar.gz", hash = "sha256:dd8d182285a0fe56bace7f45b5e7d1a6ebcbf524e8f3bd87eb0f125271b8831f"}, + {file = "mccabe-0.7.0-py2.py3-none-any.whl", hash = "sha256:6c2d30ab6be0e4a46919781807b4f0d834ebdd6c6e3dca0bda5a15f863427b6e"}, + {file = "mccabe-0.7.0.tar.gz", hash = "sha256:348e0240c33b60bbdf4e523192ef919f28cb2c3d7d5c7794f74009290f236325"}, ] [[package]] @@ -1351,13 +1352,13 @@ files = [ [[package]] name = "pycodestyle" -version = "2.7.0" +version = "2.11.1" description = "Python style guide checker" optional = false -python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*" +python-versions = ">=3.8" files = [ - {file = "pycodestyle-2.7.0-py2.py3-none-any.whl", hash = "sha256:514f76d918fcc0b55c6680472f0a37970994e07bbb80725808c17089be302068"}, - {file = "pycodestyle-2.7.0.tar.gz", hash = "sha256:c389c1d06bf7904078ca03399a4816f974a1d590090fecea0c63ec26ebaf1cef"}, + {file = "pycodestyle-2.11.1-py2.py3-none-any.whl", hash = "sha256:44fe31000b2d866f2e41841b18528a505fbd7fef9017b04eff4e2648a0fadc67"}, + {file = "pycodestyle-2.11.1.tar.gz", hash = "sha256:41ba0e7afc9752dfb53ced5489e89f8186be00e599e712660695b7a75ff2663f"}, ] [[package]] @@ -1389,13 +1390,13 @@ test = ["pytest", "pytest-cov", "pytest-regressions"] [[package]] name = "pyflakes" -version = "2.3.1" +version = "3.2.0" description = "passive checker of Python programs" optional = false -python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*" +python-versions = ">=3.8" files = [ - {file = "pyflakes-2.3.1-py2.py3-none-any.whl", hash = "sha256:7893783d01b8a89811dd72d7dfd4d84ff098e5eed95cfa8905b22bbffe52efc3"}, - {file = "pyflakes-2.3.1.tar.gz", hash = "sha256:f5bc8ecabc05bb9d291eb5203d6810b49040f6ff446a756326104746cc00c1db"}, + {file = "pyflakes-3.2.0-py2.py3-none-any.whl", hash = "sha256:84b5be138a2dfbb40689ca07e2152deb896a65c3a3e24c251c5c62489568074a"}, + {file = "pyflakes-3.2.0.tar.gz", hash = "sha256:1c61603ff154621fb2a9172037d84dca3500def8c8b630657d1701f026f8af3f"}, ] [[package]] @@ -1552,6 +1553,7 @@ files = [ {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:69b023b2b4daa7548bcfbd4aa3da05b3a74b772db9e23b982788168117739938"}, {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:81e0b275a9ecc9c0c0c07b4b90ba548307583c125f54d5b6946cfee6360c733d"}, {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba336e390cd8e4d1739f42dfe9bb83a3cc2e80f567d8805e11b46f4a943f5515"}, + {file = "PyYAML-6.0.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:326c013efe8048858a6d312ddd31d56e468118ad4cdeda36c719bf5bb6192290"}, {file = "PyYAML-6.0.1-cp310-cp310-win32.whl", hash = "sha256:bd4af7373a854424dabd882decdc5579653d7868b8fb26dc7d0e99f823aa5924"}, {file = "PyYAML-6.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:fd1592b3fdf65fff2ad0004b5e363300ef59ced41c2e6b3a99d4089fa8c5435d"}, {file = "PyYAML-6.0.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:6965a7bc3cf88e5a1c3bd2e0b5c22f8d677dc88a455344035f03399034eb3007"}, @@ -1559,8 +1561,16 @@ files = [ {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:42f8152b8dbc4fe7d96729ec2b99c7097d656dc1213a3229ca5383f973a5ed6d"}, {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:062582fca9fabdd2c8b54a3ef1c978d786e0f6b3a1510e0ac93ef59e0ddae2bc"}, {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d2b04aac4d386b172d5b9692e2d2da8de7bfb6c387fa4f801fbf6fb2e6ba4673"}, + {file = "PyYAML-6.0.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e7d73685e87afe9f3b36c799222440d6cf362062f78be1013661b00c5c6f678b"}, {file = "PyYAML-6.0.1-cp311-cp311-win32.whl", hash = "sha256:1635fd110e8d85d55237ab316b5b011de701ea0f29d07611174a1b42f1444741"}, {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"}, + {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"}, + {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"}, + {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"}, + {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"}, + {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"}, + {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"}, + {file = "PyYAML-6.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:0d3304d8c0adc42be59c5f8a4d9e3d7379e6955ad754aa9d6ab7a398b59dd1df"}, {file = "PyYAML-6.0.1-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:50550eb667afee136e9a77d6dc71ae76a44df8b3e51e41b77f6de2932bfe0f47"}, {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1fe35611261b29bd1de0070f0b2f47cb6ff71fa6595c077e42bd0c419fa27b98"}, {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:704219a11b772aea0d8ecd7058d0082713c3562b4e271b849ad7dc4a5c90c13c"}, @@ -1577,6 +1587,7 @@ files = [ {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a0cd17c15d3bb3fa06978b4e8958dcdc6e0174ccea823003a106c7d4d7899ac5"}, {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:28c119d996beec18c05208a8bd78cbe4007878c6dd15091efb73a30e90539696"}, {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7e07cbde391ba96ab58e532ff4803f79c4129397514e1413a7dc761ccd755735"}, + {file = "PyYAML-6.0.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:49a183be227561de579b4a36efbb21b3eab9651dd81b1858589f796549873dd6"}, {file = "PyYAML-6.0.1-cp38-cp38-win32.whl", hash = "sha256:184c5108a2aca3c5b3d3bf9395d50893a7ab82a38004c8f61c258d4428e80206"}, {file = "PyYAML-6.0.1-cp38-cp38-win_amd64.whl", hash = "sha256:1e2722cc9fbb45d9b87631ac70924c11d3a401b2d7f410cc0e3bbf249f2dca62"}, {file = "PyYAML-6.0.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:9eb6caa9a297fc2c2fb8862bc5370d0303ddba53ba97e71f08023b6cd73d16a8"}, @@ -1584,6 +1595,7 @@ files = [ {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5773183b6446b2c99bb77e77595dd486303b4faab2b086e7b17bc6bef28865f6"}, {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b786eecbdf8499b9ca1d697215862083bd6d2a99965554781d0d8d1ad31e13a0"}, {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bc1bf2925a1ecd43da378f4db9e4f799775d6367bdb94671027b73b393a7c42c"}, + {file = "PyYAML-6.0.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:04ac92ad1925b2cff1db0cfebffb6ffc43457495c9b3c39d3fcae417d7125dc5"}, {file = "PyYAML-6.0.1-cp39-cp39-win32.whl", hash = "sha256:faca3bdcf85b2fc05d06ff3fbc1f83e1391b3e724afa3feba7d13eeab355484c"}, {file = "PyYAML-6.0.1-cp39-cp39-win_amd64.whl", hash = "sha256:510c9deebc5c0225e8c96813043e62b680ba2f9c50a08d3724c7f28a747d1486"}, {file = "PyYAML-6.0.1.tar.gz", hash = "sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43"}, @@ -2034,4 +2046,4 @@ examples = ["cma"] [metadata] lock-version = "2.0" python-versions = ">=3.10,<3.12" -content-hash = "d772fc54f4a1b5816fcb4c6fbda968ac724b5884433f8cf604abf723830176fe" +content-hash = "07105447f85d22c61c0766e22959a4f3428f9cfa20e7499aae4017fa993fdca3" diff --git a/pyproject.toml b/pyproject.toml index d8988a3c5..09a13f1ec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,7 @@ docutils = {version = "^0.18", optional = true} cma = {version = "^3.2.2", optional = true} mypy = "^1.10.0" mypy-extensions = "^1.0.0" +flake8 = "^7.0.0" [tool.poetry.dev-dependencies] black = "24.3.0" @@ -62,10 +63,9 @@ coverage = "^6.3.3" pre-commit = "^2.19.0" pytest-html = "^3.1.1" pytest-cov = "^3.0.0" -flake8 = "^3.8" codecov = "2.1.13" click = "8.0.0" -autoflake8 = "^0.4" +autoflake = "^2.3.1" [tool.poetry.extras] docs = [ @@ -102,10 +102,15 @@ exclude = ''' )/ ''' +[tool.autoflake] +ignore-init-module-imports = true +ignore-pass-statements = true +ignore-pass-after-docstring = true + [tool.pytest.ini_options] # https://docs.pytest.org/en/6.2.x/customize.html#pyproject-toml # Directories that are not visited by pytest collector: -norecursedirs =["hooks", "*.egg", ".eggs", "dist", "build", "docs", ".tox", ".git", "__pycache__"] +norecursedirs = ["hooks", "*.egg", ".eggs", "dist", "build", "docs", ".tox", ".git", "__pycache__"] [tool.mypy] # https://mypy.readthedocs.io/en/latest/config_file.html#using-a-pyproject-toml-file From 3a06cb49d987ee84ad9794f6d298b8ca73b13bae Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sat, 4 May 2024 20:41:16 -0500 Subject: [PATCH 035/134] use isinstance --- elastica/modules/memory_block.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/elastica/modules/memory_block.py b/elastica/modules/memory_block.py index 21cb9b50f..6139e7eee 100644 --- a/elastica/modules/memory_block.py +++ b/elastica/modules/memory_block.py @@ -27,15 +27,15 @@ def construct_memory_block_structures(systems): for system_idx, sys_to_be_added in enumerate(systems): - if issubclass(sys_to_be_added.__class__, RodBase): + if isinstance(sys_to_be_added, RodBase): temp_list_for_cosserat_rod_systems.append(sys_to_be_added) temp_list_for_cosserat_rod_systems_idx.append(system_idx) - elif issubclass(sys_to_be_added.__class__, RigidBodyBase): + elif isinstance(sys_to_be_added, RigidBodyBase): temp_list_for_rigid_body_systems.append(sys_to_be_added) temp_list_for_rigid_body_systems_idx.append(system_idx) - elif issubclass(sys_to_be_added.__class__, SurfaceBase): + elif isinstance(sys_to_be_added, SurfaceBase): pass else: From 408d0bd3c1d1bf85a4d764a84faf858ff1077488 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 6 May 2024 21:22:40 +0000 Subject: [PATCH 036/134] build(deps): bump jinja2 from 3.1.3 to 3.1.4 Bumps [jinja2](https://github.com/pallets/jinja) from 3.1.3 to 3.1.4. - [Release notes](https://github.com/pallets/jinja/releases) - [Changelog](https://github.com/pallets/jinja/blob/main/CHANGES.rst) - [Commits](https://github.com/pallets/jinja/compare/3.1.3...3.1.4) --- updated-dependencies: - dependency-name: jinja2 dependency-type: indirect ... Signed-off-by: dependabot[bot] --- poetry.lock | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/poetry.lock b/poetry.lock index 4c0f39226..012ac94f5 100644 --- a/poetry.lock +++ b/poetry.lock @@ -604,13 +604,13 @@ files = [ [[package]] name = "jinja2" -version = "3.1.3" +version = "3.1.4" description = "A very fast and expressive template engine." -optional = true +optional = false python-versions = ">=3.7" files = [ - {file = "Jinja2-3.1.3-py3-none-any.whl", hash = "sha256:7d6d50dd97d52cbc355597bd845fabfbac3f551e1f99619e39a35ce8c370b5fa"}, - {file = "Jinja2-3.1.3.tar.gz", hash = "sha256:ac8bd6544d4bb2c9792bf3a159e80bba8fda7f07e81bc3aed565432d5925ba90"}, + {file = "jinja2-3.1.4-py3-none-any.whl", hash = "sha256:bc5dd2abb727a5319567b7a813e6a2e7318c39f4f487cfe6c89c6f9c7d25197d"}, + {file = "jinja2-3.1.4.tar.gz", hash = "sha256:4a3aee7acbbe7303aede8e9648d13b8bf88a429282aa6122a993f0ac800cb369"}, ] [package.dependencies] @@ -793,7 +793,7 @@ testing = ["coverage", "pytest", "pytest-cov", "pytest-regressions"] name = "markupsafe" version = "2.1.5" description = "Safely add untrusted strings to HTML/XML markup." -optional = true +optional = false python-versions = ">=3.7" files = [ {file = "MarkupSafe-2.1.5-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:a17a92de5231666cfbe003f0e4b9b3a7ae3afb1ec2845aadc2bacc93ff85febc"}, From 888bfbdc95509511c1069c1aea962494d8c0406b Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Tue, 7 May 2024 01:34:54 -0500 Subject: [PATCH 037/134] typing: base_system module --- elastica/modules/base_system.py | 104 ++++++++++++++++++++------------ elastica/typing.py | 9 ++- 2 files changed, 72 insertions(+), 41 deletions(-) diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index 34f42c7d4..e14d8e422 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -5,7 +5,16 @@ Basic coordinating for multiple, smaller systems that have an independently integrable interface (i.e. works with symplectic or explicit routines `timestepper.py`.) """ -from typing import Iterable, Callable, AnyStr +from typing import Iterable, Callable, AnyStr, Type, Generator +from typing import final +from elastica.typing import ( + SystemType, + SynchronizeOperator, + ConstrainValuesOperator, + ConstrainRatesOperator, + CallbackOperator, + FinalizeOperator, +) import numpy as np @@ -31,11 +40,6 @@ class BaseSystemCollection(MutableSequence): _systems: list List of rod-like objects. - """ - - """ - Developer Note - ----- Note ---- We can directly subclass a list for the @@ -47,23 +51,31 @@ def __init__(self): # Collection of functions. Each group is executed as a collection at the different steps. # Each component (Forcing, Connection, etc.) registers the executable (callable) function # in the group that that needs to be executed. These should be initialized before mixin. - self._feature_group_synchronize: Iterable[Callable[[float], None]] = [] - self._feature_group_constrain_values: Iterable[Callable[[float], None]] = [] - self._feature_group_constrain_rates: Iterable[Callable[[float], None]] = [] - self._feature_group_callback: Iterable[Callable[[float, int, AnyStr], None]] = ( - [] - ) - self._feature_group_finalize: Iterable[Callable] = [] + self._feature_group_synchronize: list[SynchronizeOperator] = [] + self._feature_group_constrain_values: list[ConstrainValuesOperator] = [] + self._feature_group_constrain_rates: list[ConstrainRatesOperator] = [] + self._feature_group_callback: list[CallbackOperator] = [] + self._feature_group_finalize: list[FinalizeOperator] = [] + # We need to initialize our mixin classes super(BaseSystemCollection, self).__init__() + # List of system types/bases that are allowed - self.allowed_sys_types = (RodBase, RigidBodyBase, SurfaceBase) + self.allowed_sys_types: tuple[Type[SystemType]] = ( + RodBase, + RigidBodyBase, + SurfaceBase, + ) + # List of systems to be integrated - self._systems = [] + self._systems: list[SystemType] = [] + self._memory_blocks: list[SystemType] = [] + # Flag Finalize: Finalizing twice will cause an error, # but the error message is very misleading - self._finalize_flag = False + self._finalize_flag: bool = False + @final def _check_type(self, sys_to_be_added: AnyStr): if not issubclass(sys_to_be_added.__class__, self.allowed_sys_types): raise TypeError( @@ -78,38 +90,40 @@ def _check_type(self, sys_to_be_added: AnyStr): ) return True - def __len__(self): + def __len__(self) -> int: return len(self._systems) - def __getitem__(self, idx): + def __getitem__(self, idx: int) -> SystemType: return self._systems[idx] - def __delitem__(self, idx): + def __delitem__(self, idx: int) -> None: del self._systems[idx] - def __setitem__(self, idx, system): + def __setitem__(self, idx: int, system: SystemType) -> None: self._check_type(system) self._systems[idx] = system - def insert(self, idx, system): + def insert(self, idx: int, system: SystemType) -> None: self._check_type(system) self._systems.insert(idx, system) - def __str__(self): + def __str__(self) -> str: + """To be readable""" return str(self._systems) - def extend_allowed_types(self, additional_types): + @final + def extend_allowed_types(self, additional_types) -> None: self.allowed_sys_types += additional_types - def override_allowed_types(self, allowed_types): + @final + def override_allowed_types(self, allowed_types) -> None: self.allowed_sys_types = allowed_types + @final def _get_sys_idx_if_valid(self, sys_to_be_added): - from numpy import int_ as npint - - n_systems = len(self._systems) # Total number of systems from mixed-in class + n_systems = len(self) # Total number of systems from mixed-in class - if isinstance(sys_to_be_added, (int, npint)): + if isinstance(sys_to_be_added, (int, np.int_)): # 1. If they are indices themselves, check range assert ( -n_systems <= sys_to_be_added < n_systems @@ -131,7 +145,14 @@ def _get_sys_idx_if_valid(self, sys_to_be_added): return sys_idx - def finalize(self): + @final + def blocks(self) -> Generator[SystemType, None, None]: + assert self._finalize_flag, "The simulator is not finalized." + for block in self._memory_blocks: + yield block + + @final + def finalize(self) -> None: """ This method finalizes the simulator class. When it is called, it is assumed that the user has appended all rod-like objects to the simulator as well as all boundary conditions, callbacks, etc., @@ -139,8 +160,8 @@ def finalize(self): the user cannot add new features to the simulator class. """ - # This generates more straight-forward error. - assert self._finalize_flag is not True, "The finalize cannot be called twice." + assert not self._finalize_flag, "The finalize cannot be called twice." + self._finalize_flag = True # construct memory block self._memory_blocks = construct_memory_block_structures(self._systems) @@ -150,13 +171,13 @@ def finalize(self): periodic boundaries, a new constrain for memory block rod added called as _ConstrainPeriodicBoundaries. This constrain will synchronize the only periodic boundaries of position, director, velocity and omega variables. """ - for i in range(len(self._memory_blocks)): + for block in self._memory_blocks: # append the memory block to the simulation as a system. Memory block is the final system in the simulation. - self.append(self._memory_blocks[i]) - if hasattr(self._memory_blocks[i], "ring_rod_flag"): + self.append(block) + if hasattr(block, "ring_rod_flag"): # Apply the constrain to synchronize the periodic boundaries of the memory rod. Find the memory block # sys idx among other systems added and then apply boundary conditions. - memory_block_idx = self._get_sys_idx_if_valid(self._memory_blocks[i]) + memory_block_idx = self._get_sys_idx_if_valid(block) self.constrain(self._systems[memory_block_idx]).using( _ConstrainPeriodicBoundaries, ) @@ -170,7 +191,6 @@ def finalize(self): self._feature_group_finalize = None # Toggle the finalize_flag - self._finalize_flag = True # sort _feature_group_synchronize so that _call_contacts is at the end _call_contacts_index = [] for idx, feature in enumerate(self._feature_group_synchronize): @@ -183,22 +203,26 @@ def finalize(self): self._feature_group_synchronize.pop(index) ) - def synchronize(self, time: np.floating): + @final + def synchronize(self, time: np.floating) -> None: # Collection call _feature_group_synchronize for feature in self._feature_group_synchronize: feature(time) - def constrain_values(self, time: np.floating): + @final + def constrain_values(self, time: np.floating) -> None: # Collection call _feature_group_constrain_values for feature in self._feature_group_constrain_values: feature(time) - def constrain_rates(self, time: np.floating): + @final + def constrain_rates(self, time: np.floating) -> None: # Collection call _feature_group_constrain_rates for feature in self._feature_group_constrain_rates: feature(time) - def apply_callbacks(self, time: np.floating, current_step: int): + @final + def apply_callbacks(self, time: np.floating, current_step: int) -> None: # Collection call _feature_group_callback for feature in self._feature_group_callback: feature(time, current_step) diff --git a/elastica/typing.py b/elastica/typing.py index b6bd654e6..3abc59e3d 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -4,7 +4,7 @@ """ from typing import TYPE_CHECKING -from typing import Type, Union, Callable, Any +from typing import Type, Union, Callable, Any, AnyStr from typing import TypeAlias @@ -77,3 +77,10 @@ RodType: TypeAlias = Type[RodBase] SystemCollectionType: TypeAlias = BaseSystemCollection AllowedContactType: TypeAlias = Union[SystemType, Type[SurfaceBase]] + +# Operators in elastica.modules +SynchronizeOperator: TypeAlias = Callable[[float], None] +ConstrainValuesOperator: TypeAlias = Callable[[float], None] +ConstrainRatesOperator: TypeAlias = Callable[[float], None] +CallbackOperator: TypeAlias = Callable[[float, int, AnyStr], None] +FinalizeOperator: TypeAlias = Callable[[], None] From ff4fbc21b20ce592015d515838bdb7cceb0363ab Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Tue, 7 May 2024 01:57:40 -0500 Subject: [PATCH 038/134] wip: callback module typing --- elastica/modules/base_system.py | 3 ++- elastica/modules/callbacks.py | 14 ++++++++------ elastica/modules/protocol.py | 1 + elastica/typing.py | 1 + 4 files changed, 12 insertions(+), 7 deletions(-) create mode 100644 elastica/modules/protocol.py diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index e14d8e422..dadffe4ae 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -9,6 +9,7 @@ from typing import final from elastica.typing import ( SystemType, + SystemIdxType, SynchronizeOperator, ConstrainValuesOperator, ConstrainRatesOperator, @@ -120,7 +121,7 @@ def override_allowed_types(self, allowed_types) -> None: self.allowed_sys_types = allowed_types @final - def _get_sys_idx_if_valid(self, sys_to_be_added): + def _get_sys_idx_if_valid(self, sys_to_be_added: SystemType) -> SystemIdxType: n_systems = len(self) # Total number of systems from mixed-in class if isinstance(sys_to_be_added, (int, np.int_)): diff --git a/elastica/modules/callbacks.py b/elastica/modules/callbacks.py index af1482add..c0a654d8e 100644 --- a/elastica/modules/callbacks.py +++ b/elastica/modules/callbacks.py @@ -4,8 +4,10 @@ Provides the callBack interface to collect data over time (see `callback_functions.py`). """ +from elastica.typing import SystemType, SystemIdxType from elastica.callback_functions import CallBackBaseClass +from .base_system import BaseSystemCollection class CallBacks: @@ -20,13 +22,13 @@ class CallBacks: List of call back classes defined for rod-like objects. """ - def __init__(self): - self._callback_list = [] + def __init__(self: BaseSystemCollection): + self._callback_list: list[_Callback] = [] super(CallBacks, self).__init__() self._feature_group_callback.append(self._callback_execution) self._feature_group_finalize.append(self._finalize_callback) - def collect_diagnostics(self, system): + def collect_diagnostics(self: BaseSystemCollection, system: SystemType): """ This method calls user-defined call-back classes for a user-defined system or rod-like object. You need to input the @@ -41,10 +43,10 @@ def collect_diagnostics(self, system): ------- """ - sys_idx = self._get_sys_idx_if_valid(system) + sys_idx: SystemIdxType = self._get_sys_idx_if_valid(system) # Create _Constraint object, cache it and return to user - _callbacks = _CallBack(sys_idx) + _callbacks: _Callback = _CallBack(sys_idx) self._callback_list.append(_callbacks) return _callbacks @@ -92,7 +94,7 @@ class _CallBack: Arbitrary keyword arguments. """ - def __init__(self, sys_idx: int): + def __init__(self, sys_idx: SystemIdxType): """ Parameters diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py new file mode 100644 index 000000000..9253c4dff --- /dev/null +++ b/elastica/modules/protocol.py @@ -0,0 +1 @@ +from typing import Protocol diff --git a/elastica/typing.py b/elastica/typing.py index 3abc59e3d..acc0cbf4d 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -33,6 +33,7 @@ SystemType: TypeAlias = Union[SymplecticSystemProtocol, ExplicitSystemProtocol] +SystemIdxType: TypeAlias = int # TODO: Modify this line and move to elastica/typing.py once system state is defined # Mostly used in explicit stepper: for symplectic, use kinetic and dynamic state From df61998c96d31c631316ef583bd88a3278d38144 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sat, 11 May 2024 00:37:38 -0500 Subject: [PATCH 039/134] typing: callback module typing New protocol for SystemCollection --- elastica/callback_functions.py | 2 +- elastica/modules/base_system.py | 6 +- elastica/modules/callbacks.py | 95 ++++++++++++++++++---------- elastica/modules/operator_group.py | 3 +- elastica/modules/protocol.py | 42 ++++++++++++ elastica/typing.py | 32 ++++++++-- tests/test_modules/test_callbacks.py | 2 +- 7 files changed, 136 insertions(+), 46 deletions(-) diff --git a/elastica/callback_functions.py b/elastica/callback_functions.py index ab865dfd0..9214e00cb 100644 --- a/elastica/callback_functions.py +++ b/elastica/callback_functions.py @@ -25,7 +25,7 @@ def __init__(self): """ pass - def make_callback(self, system, time, current_step: int): + def make_callback(self, system: SystemType, time: np.floating, current_step: int): """ This method is called every time step. Users can define which parameters are called back and recorded. Also users diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index 7de33b05d..56b350cc5 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -5,14 +5,14 @@ Basic coordinating for multiple, smaller systems that have an independently integrable interface (i.e. works with symplectic or explicit routines `timestepper.py`.) """ -from typing import Iterable, Callable, AnyStr, Type, Generator +from typing import Iterable, AnyStr, Type, Generator from typing import final from elastica.typing import ( SystemType, SystemIdxType, OperatorType, OperatorCallbackType, - OperatorFinalizeType + OperatorFinalizeType, ) import numpy as np @@ -61,7 +61,7 @@ def __init__(self): self._feature_group_callback: Iterable[OperatorCallbackType] = [] self._feature_group_finalize: Iterable[OperatorFinalizeType] = [] # We need to initialize our mixin classes - super(BaseSystemCollection, self).__init__() + super().__init__() # List of system types/bases that are allowed self.allowed_sys_types: tuple[Type[SystemType]] = ( diff --git a/elastica/modules/callbacks.py b/elastica/modules/callbacks.py index c0a654d8e..d6906fb08 100644 --- a/elastica/modules/callbacks.py +++ b/elastica/modules/callbacks.py @@ -4,10 +4,33 @@ Provides the callBack interface to collect data over time (see `callback_functions.py`). """ -from elastica.typing import SystemType, SystemIdxType +from typing import Type, Protocol +from typing_extensions import Self # 3.11: from typing import Self +from elastica.typing import SystemType, SystemIdxType, OperatorFinalizeType +from elastica.typing import CallbackParam +from .protocol import ModuleProtocol + +import numpy as np from elastica.callback_functions import CallBackBaseClass -from .base_system import BaseSystemCollection +from .protocol import SystemCollectionProtocol + + +class SystemCollectionWithCallBacksProtocol(SystemCollectionProtocol, Protocol): + _callback_list: list[ModuleProtocol] + _callback_operators: list[tuple[int, CallBackBaseClass]] + + _finalize_callback: OperatorFinalizeType + + def collect_diagnostics(self, system: SystemType) -> ModuleProtocol: ... + + def _callback_execution( + self, + time: np.floating, + current_step: int, + *args: CallbackParam.args, + **kwargs: CallbackParam.kwargs, + ) -> None: ... class CallBacks: @@ -22,13 +45,16 @@ class CallBacks: List of call back classes defined for rod-like objects. """ - def __init__(self: BaseSystemCollection): - self._callback_list: list[_Callback] = [] + def __init__(self: SystemCollectionWithCallBacksProtocol) -> None: + self._callback_list: list[ModuleProtocol] = [] + self._callback_operators: list[tuple[int, CallBackBaseClass]] = [] super(CallBacks, self).__init__() self._feature_group_callback.append(self._callback_execution) self._feature_group_finalize.append(self._finalize_callback) - def collect_diagnostics(self: BaseSystemCollection, system: SystemType): + def collect_diagnostics( + self: SystemCollectionWithCallBacksProtocol, system: SystemType + ) -> ModuleProtocol: """ This method calls user-defined call-back classes for a user-defined system or rod-like object. You need to input the @@ -46,35 +72,29 @@ def collect_diagnostics(self: BaseSystemCollection, system: SystemType): sys_idx: SystemIdxType = self._get_sys_idx_if_valid(system) # Create _Constraint object, cache it and return to user - _callbacks: _Callback = _CallBack(sys_idx) + _callbacks: ModuleProtocol = _CallBack(sys_idx) self._callback_list.append(_callbacks) return _callbacks - def _finalize_callback(self): - # From stored _CallBack objects, instantiate the boundary conditions - # inplace : https://stackoverflow.com/a/1208792 - + def _finalize_callback(self: SystemCollectionWithCallBacksProtocol) -> None: # dev : the first index stores the rod index to collect data. - # Technically we can use another array but it its one more book-keeping - # step. Being lazy, I put them both in the same array - self._callback_list[:] = [ - (callback.id(), callback(self._systems[callback.id()])) - for callback in self._callback_list + self._callback_operators = [ + (callback.id(), callback.instantiate()) for callback in self._callback_list ] - # Sort from lowest id to highest id for potentially better memory access - # _callbacks contains list of tuples. First element of tuple is rod number and - # following elements are the type of boundary condition such as - # [(0, MyCallBack), (1, MyVelocityCallBack), ... ] - # Thus using lambda we iterate over the list of tuples and use rod number (x[0]) - # to sort callbacks. - self._callback_list.sort(key=lambda x: x[0]) - - self._callback_execution(time=0.0, current_step=0) - - def _callback_execution(self, time, current_step: int, *args, **kwargs): - for sys_id, callback in self._callback_list: + # First callback execution + time = np.float64(0.0) + self._callback_execution(time=time, current_step=0) + + def _callback_execution( + self: SystemCollectionWithCallBacksProtocol, + time: np.floating, + current_step: int, + *args: CallbackParam.args, + **kwargs: CallbackParam.kwargs, + ) -> None: + for sys_id, callback in self._callback_operators: callback.make_callback( self._systems[sys_id], time, current_step, *args, **kwargs ) @@ -102,12 +122,17 @@ def __init__(self, sys_idx: SystemIdxType): sys_idx: int rod object index """ - self._sys_idx = sys_idx - self._callback_cls = None - self._args = () - self._kwargs = {} - - def using(self, callback_cls, *args, **kwargs): + self._sys_idx: SystemIdxType = sys_idx + self._callback_cls: Type[CallBackBaseClass] + self._args: CallbackParam.args + self._kwargs: CallbackParam.kwargs + + def using( + self, + callback_cls: Type[CallBackBaseClass], + *args: CallbackParam.args, + **kwargs: CallbackParam.kwargs, + ) -> Self: """ This method is a module to set which callback class is used to collect data from user defined rod-like object. @@ -131,10 +156,10 @@ def using(self, callback_cls, *args, **kwargs): self._kwargs = kwargs return self - def id(self): + def id(self) -> SystemIdxType: return self._sys_idx - def __call__(self, *args, **kwargs) -> CallBackBaseClass: + def instantiate(self) -> CallBackBaseClass: """Constructs a callback functions after checks""" if not self._callback_cls: raise RuntimeError( diff --git a/elastica/modules/operator_group.py b/elastica/modules/operator_group.py index e18b12322..96d14d900 100644 --- a/elastica/modules/operator_group.py +++ b/elastica/modules/operator_group.py @@ -1,4 +1,4 @@ -from typing import Generic, TypeVar +from typing import TypeVar from collections.abc import Iterable @@ -7,6 +7,7 @@ T = TypeVar("T") F = TypeVar("F") + class OperatorGroupFIFO(Iterable): """ A class to store the features and their corresponding operators in a FIFO manner. diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index 9253c4dff..2145c726a 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -1 +1,43 @@ from typing import Protocol +from typing_extensions import Self # 3.11: from typing import Self + +from elastica.typing import ( + SystemIdxType, + OperatorType, + OperatorCallbackType, + OperatorFinalizeType, + SystemType, +) + +from collections.abc import MutableSequence + + +class SystemCollectionProtocol(Protocol, MutableSequence): + _systems: list[SystemType] + + @property + def _feature_group_synchronize(self) -> Iterable[OperatorType]: ... + + @property + def _feature_group_constrain_values(self) -> Iterable[OperatorType]: ... + + @property + def _feature_group_constrain_rates(self) -> Iterable[OperatorType]: ... + + @property + def _feature_group_callback(self) -> Iterable[OperatorCallbackType]: ... + + @property + def _feature_group_finalize(self) -> Iterable[OperatorFinalizeType]: ... + + def blocks(self) -> Generator[SystemType, None, None]: ... + + def _get_sys_idx_if_valid(self, sys_to_be_added: SystemType) -> SystemIdxType: ... + + +class ModuleProtocol(Protocol): + def using(self, callback_cls, *args, **kwargs) -> Self: ... + + def instantiate(self) -> Type: ... + + def id(self) -> SystemIdxType: ... diff --git a/elastica/typing.py b/elastica/typing.py index 6bbb27792..6c7961303 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -4,7 +4,7 @@ """ from typing import TYPE_CHECKING -from typing import Type, Union, Callable, Any, AnyStr +from typing import Type, Union, Callable, Any, ParamSpec from typing import TypeAlias @@ -14,8 +14,8 @@ from .rod import RodBase from .rigidbody import RigidBodyBase from .surface import SurfaceBase - from .modules import BaseSystemCollection + from .modules.protocol import SystemCollectionProtocol from .rod.data_structures import _State as State from .systems.protocol import SymplecticSystemProtocol, ExplicitSystemProtocol from .timestepper.protocol import ( @@ -23,11 +23,20 @@ SymplecticStepperProtocol, MemoryProtocol, ) + + # Modules Base Classes + from .boundary_conditions import FreeBC + from .callback_functions import CallBackBaseClass + from .contact_forces import NoContact + from .dissipation import DamperBase + from .external_forces import NoForces + from .joint import FreeJoint else: RodBase = None RigidBodyBase = None SurfaceBase = None - BaseSystemCollection = None + + SystemCollectionProtocol = None State = "State" SymplecticSystemProtocol = None @@ -37,10 +46,22 @@ SymplecticStepperProtocol = None MemoryProtocol = None + # Modules Base Classes + FreeBC = None + CallBackBaseClass = None + NoContact = None + DamperBase = None + NoForces = None + FreeJoint = None + SystemType: TypeAlias = Union[SymplecticSystemProtocol, ExplicitSystemProtocol] SystemIdxType: TypeAlias = int +ModuleObjectTypes: TypeAlias = ( + NoForces | NoContact | FreeJoint | FreeBC | DamperBase | CallBackBaseClass +) + # TODO: Modify this line and move to elastica/typing.py once system state is defined # Mostly used in explicit stepper: for symplectic, use kinetic and dynamic state StateType: TypeAlias = State @@ -81,10 +102,11 @@ # ] RodType: TypeAlias = Type[RodBase] -SystemCollectionType: TypeAlias = BaseSystemCollection +SystemCollectionType: TypeAlias = SystemCollectionProtocol AllowedContactType: TypeAlias = Union[SystemType, Type[SurfaceBase]] # Operators in elastica.modules +CallbackParam = ParamSpec("CallbackParam") OperatorType: TypeAlias = Callable[[float], None] -OperatorCallbackType: TypeAlias = Callable[[float, int], None] +OperatorCallbackType: TypeAlias = Callable[CallbackParam, None] OperatorFinalizeType: TypeAlias = Callable diff --git a/tests/test_modules/test_callbacks.py b/tests/test_modules/test_callbacks.py index 57d4a4268..e9c04f57e 100644 --- a/tests/test_modules/test_callbacks.py +++ b/tests/test_modules/test_callbacks.py @@ -55,7 +55,7 @@ def mock_init(self, *args, **kwargs): # Actual test is here, this should not throw with pytest.raises(TypeError) as excinfo: - _ = callback() + _ = callback.instantiate() assert "Unable to construct" in str(excinfo.value) From 1994800097994df4532ca0ff87d8e9d6f1cfbb4c Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Mon, 13 May 2024 00:15:26 -0500 Subject: [PATCH 040/134] test: fix callback module related tests --- elastica/callback_functions.py | 2 ++ elastica/typing.py | 6 +++--- tests/test_modules/test_base_system.py | 2 +- tests/test_modules/test_callbacks.py | 4 +++- 4 files changed, 9 insertions(+), 5 deletions(-) diff --git a/elastica/callback_functions.py b/elastica/callback_functions.py index 9214e00cb..7760694ea 100644 --- a/elastica/callback_functions.py +++ b/elastica/callback_functions.py @@ -1,10 +1,12 @@ __doc__ = """ Module contains callback classes to save simulation data for rod-like objects """ +from elastica.typing import SystemType import os import sys import numpy as np import logging + from collections import defaultdict diff --git a/elastica/typing.py b/elastica/typing.py index 6c7961303..377b98a62 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -58,9 +58,9 @@ SystemType: TypeAlias = Union[SymplecticSystemProtocol, ExplicitSystemProtocol] SystemIdxType: TypeAlias = int -ModuleObjectTypes: TypeAlias = ( - NoForces | NoContact | FreeJoint | FreeBC | DamperBase | CallBackBaseClass -) +# ModuleObjectTypes: TypeAlias = ( +# NoForces | NoContact | FreeJoint | FreeBC | DamperBase | CallBackBaseClass +# ) # TODO: Modify this line and move to elastica/typing.py once system state is defined # Mostly used in explicit stepper: for symplectic, use kinetic and dynamic state diff --git a/tests/test_modules/test_base_system.py b/tests/test_modules/test_base_system.py index 76f63cd4e..e26791d66 100644 --- a/tests/test_modules/test_base_system.py +++ b/tests/test_modules/test_base_system.py @@ -221,7 +221,7 @@ def test_callback(self, load_collection, legal_callback): simulator_class.collect_diagnostics(rod).using(legal_callback) simulator_class.finalize() # After finalize check if the created callback object is instance of the class we have given. - assert isinstance(simulator_class._callback_list[-1][-1], legal_callback) + assert isinstance(simulator_class._callback_operators[-1][-1], legal_callback) # TODO: this is a dummy test for apply_callbacks find a better way to test them simulator_class.apply_callbacks(time=0, current_step=0) diff --git a/tests/test_modules/test_callbacks.py b/tests/test_modules/test_callbacks.py index e9c04f57e..b6216cc5f 100644 --- a/tests/test_modules/test_callbacks.py +++ b/tests/test_modules/test_callbacks.py @@ -164,10 +164,12 @@ def test_callback_finalize_correctness(self, load_rod_with_callbacks): scwc._finalize_callback() - for x, y in scwc._callback_list: + for x, y in scwc._callback_operators: assert type(x) is int assert type(y) is callback_cls + assert not hasattr(scwc, "_callback_list") + @pytest.mark.xfail def test_callback_finalize_sorted(self, load_rod_with_callbacks): scwc, callback_cls = load_rod_with_callbacks From 07dd6f47a4637fafaafbed75f77f826922af5e72 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Mon, 13 May 2024 00:55:59 -0500 Subject: [PATCH 041/134] typing: base system with requisite as part of system protocol --- elastica/modules/base_system.py | 37 +++++++++++++++++---------------- elastica/modules/callbacks.py | 8 +++---- elastica/modules/protocol.py | 6 +++--- elastica/systems/protocol.py | 7 ++++--- 4 files changed, 29 insertions(+), 29 deletions(-) diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index 56b350cc5..900239c69 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -5,7 +5,7 @@ Basic coordinating for multiple, smaller systems that have an independently integrable interface (i.e. works with symplectic or explicit routines `timestepper.py`.) """ -from typing import Iterable, AnyStr, Type, Generator +from typing import Type, Generator, Iterable from typing import final from elastica.typing import ( SystemType, @@ -51,20 +51,20 @@ class BaseSystemCollection(MutableSequence): https://stackoverflow.com/q/3945940 """ - def __init__(self): + def __init__(self) -> None: # Collection of functions. Each group is executed as a collection at the different steps. # Each component (Forcing, Connection, etc.) registers the executable (callable) function # in the group that that needs to be executed. These should be initialized before mixin. self._feature_group_synchronize: Iterable[OperatorType] = OperatorGroupFIFO() - self._feature_group_constrain_values: Iterable[OperatorType] = [] - self._feature_group_constrain_rates: Iterable[OperatorType] = [] - self._feature_group_callback: Iterable[OperatorCallbackType] = [] - self._feature_group_finalize: Iterable[OperatorFinalizeType] = [] + self._feature_group_constrain_values: list[OperatorType] = [] + self._feature_group_constrain_rates: list[OperatorType] = [] + self._feature_group_callback: list[OperatorCallbackType] = [] + self._feature_group_finalize: list[OperatorFinalizeType] = [] # We need to initialize our mixin classes super().__init__() # List of system types/bases that are allowed - self.allowed_sys_types: tuple[Type[SystemType]] = ( + self.allowed_sys_types: tuple[Type[SystemType], ...] = ( RodBase, RigidBodyBase, SurfaceBase, @@ -79,8 +79,8 @@ def __init__(self): self._finalize_flag: bool = False @final - def _check_type(self, sys_to_be_added: AnyStr): - if not issubclass(sys_to_be_added.__class__, self.allowed_sys_types): + def _check_type(self, sys_to_be_added: SystemType) -> bool: + if not isinstance(sys_to_be_added, self.allowed_sys_types): raise TypeError( "{0}\n" "is not a system passing validity\n" @@ -96,17 +96,17 @@ def _check_type(self, sys_to_be_added: AnyStr): def __len__(self) -> int: return len(self._systems) - def __getitem__(self, idx: int) -> SystemType: + def __getitem__(self, idx: int | slice, /) -> SystemType | list[SystemType]: # type: ignore return self._systems[idx] - def __delitem__(self, idx: int) -> None: + def __delitem__(self, idx: int | slice, /) -> None: # type: ignore del self._systems[idx] - def __setitem__(self, idx: int, system: SystemType) -> None: + def __setitem__(self, idx: int | slice, system: SystemType | Iterable[SystemType], /) -> None: # type: ignore self._check_type(system) self._systems[idx] = system - def insert(self, idx: int, system: SystemType) -> None: + def insert(self, idx: int, system: SystemType) -> None: # type: ignore self._check_type(system) self._systems.insert(idx, system) @@ -116,20 +116,21 @@ def __str__(self) -> str: @final def extend_allowed_types( - self, additional_types: list[Type[SystemType], ...] + self, additional_types: tuple[Type[SystemType], ...] ) -> None: self.allowed_sys_types += additional_types @final def override_allowed_types( - self, allowed_types: list[Type[SystemType], ...] + self, allowed_types: tuple[Type[SystemType], ...] ) -> None: - self.allowed_sys_types = tuple(allowed_types) + self.allowed_sys_types = allowed_types @final def _get_sys_idx_if_valid(self, sys_to_be_added: SystemType) -> SystemIdxType: n_systems = len(self) # Total number of systems from mixed-in class + sys_idx: SystemIdxType if isinstance(sys_to_be_added, (int, np.int_)): # 1. If they are indices themselves, check range assert ( @@ -137,7 +138,7 @@ def _get_sys_idx_if_valid(self, sys_to_be_added: SystemType) -> SystemIdxType: ), "Rod index {} exceeds number of registered rodtems".format( sys_to_be_added ) - sys_idx = sys_to_be_added + sys_idx = int(sys_to_be_added) elif self._check_type(sys_to_be_added): # 2. If they are rod objects (most likely), lookup indices # index might have some problems : https://stackoverflow.com/a/176921 @@ -195,7 +196,7 @@ def finalize(self) -> None: # Clear the finalize feature group, just for the safety. self._feature_group_finalize.clear() - self._feature_group_finalize = None + del self._feature_group_finalize @final def synchronize(self, time: np.floating) -> None: diff --git a/elastica/modules/callbacks.py b/elastica/modules/callbacks.py index d6906fb08..472b042c4 100644 --- a/elastica/modules/callbacks.py +++ b/elastica/modules/callbacks.py @@ -82,6 +82,8 @@ def _finalize_callback(self: SystemCollectionWithCallBacksProtocol) -> None: self._callback_operators = [ (callback.id(), callback.instantiate()) for callback in self._callback_list ] + self._callback_list.clear() + del self._callback_list # First callback execution time = np.float64(0.0) @@ -91,13 +93,9 @@ def _callback_execution( self: SystemCollectionWithCallBacksProtocol, time: np.floating, current_step: int, - *args: CallbackParam.args, - **kwargs: CallbackParam.kwargs, ) -> None: for sys_id, callback in self._callback_operators: - callback.make_callback( - self._systems[sys_id], time, current_step, *args, **kwargs - ) + callback.make_callback(self._systems[sys_id], time, current_step) class _CallBack: diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index 2145c726a..6007eb454 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -1,4 +1,4 @@ -from typing import Protocol +from typing import Protocol, Type, Generator, Iterable from typing_extensions import Self # 3.11: from typing import Self from elastica.typing import ( @@ -9,10 +9,10 @@ SystemType, ) -from collections.abc import MutableSequence +pass -class SystemCollectionProtocol(Protocol, MutableSequence): +class SystemCollectionProtocol(Protocol): _systems: list[SystemType] @property diff --git a/elastica/systems/protocol.py b/elastica/systems/protocol.py index 3d78f32e0..1661a87d2 100644 --- a/elastica/systems/protocol.py +++ b/elastica/systems/protocol.py @@ -1,8 +1,7 @@ __doc__ = """Base class for elastica system""" -from typing import Protocol - -from elastica.typing import StateType +from typing import Protocol, Type +from elastica.typing import StateType, SystemType from elastica.rod.data_structures import _KinematicState, _DynamicState import numpy as np @@ -14,6 +13,8 @@ class SystemProtocol(Protocol): Protocol for all elastica system """ + REQUISITE_MODULES: list[Type[SystemType]] + @property def n_nodes(self) -> int: ... From 4279c68e7fc827c4bf25062f58491a82188c7cdc Mon Sep 17 00:00:00 2001 From: armantekinalp Date: Mon, 13 May 2024 14:38:27 -0500 Subject: [PATCH 042/134] memory block typing for rod, rigid body and rod base --- .../memory_block/memory_block_rigid_body.py | 14 +++++----- elastica/memory_block/memory_block_rod.py | 18 +++++++------ .../memory_block/memory_block_rod_base.py | 27 ++++++++++--------- 3 files changed, 33 insertions(+), 26 deletions(-) diff --git a/elastica/memory_block/memory_block_rigid_body.py b/elastica/memory_block/memory_block_rigid_body.py index 827363784..9891da7b2 100644 --- a/elastica/memory_block/memory_block_rigid_body.py +++ b/elastica/memory_block/memory_block_rigid_body.py @@ -7,7 +7,7 @@ class MemoryBlockRigidBody(RigidBodyBase, _RigidRodSymplecticStepperMixin): - def __init__(self, systems: Sequence, system_idx_list: Sequence[np.int64]): + def __init__(self, systems: Sequence, system_idx_list: Sequence[np.int64]) -> None: self.n_bodies = len(systems) self.n_elems = self.n_bodies @@ -23,7 +23,7 @@ def __init__(self, systems: Sequence, system_idx_list: Sequence[np.int64]): # Initialize the mixin class for symplectic time-stepper. _RigidRodSymplecticStepperMixin.__init__(self) - def _allocate_block_variables_scalars(self, systems: Sequence): + def _allocate_block_variables_scalars(self, systems: Sequence) -> None: """ This function takes system collection and allocates the variables for block-structure and references allocated variables back to the systems. @@ -62,7 +62,7 @@ def _allocate_block_variables_scalars(self, systems: Sequence): value_type="scalar", ) - def _allocate_block_variables_vectors(self, systems: Sequence): + def _allocate_block_variables_vectors(self, systems: Sequence) -> None: """ This function takes system collection and allocates the vector variables for block-structure and references allocated vector variables back to the systems. @@ -98,7 +98,7 @@ def _allocate_block_variables_vectors(self, systems: Sequence): value_type="vector", ) - def _allocate_block_variables_matrix(self, systems: Sequence): + def _allocate_block_variables_matrix(self, systems: Sequence) -> None: """ This function takes system collection and allocates the matrix variables for block-structure and references allocated matrix variables back to the systems. @@ -134,7 +134,9 @@ def _allocate_block_variables_matrix(self, systems: Sequence): value_type="tensor", ) - def _allocate_block_variables_for_symplectic_stepper(self, systems: Sequence): + def _allocate_block_variables_for_symplectic_stepper( + self, systems: Sequence + ) -> None: """ This function takes system collection and allocates the variables used by symplectic stepper for block-structure and references allocated variables back to the systems. @@ -201,7 +203,7 @@ def _map_system_properties_to_block_memory( """ if value_type == "scalar": - view_shape = (self.n_elems,) + view_shape: tuple[int, ...] = (self.n_elems,) elif value_type == "vector": view_shape = (3, self.n_elems) diff --git a/elastica/memory_block/memory_block_rod.py b/elastica/memory_block/memory_block_rod.py index b6a6b73f7..ad2ed0bbb 100644 --- a/elastica/memory_block/memory_block_rod.py +++ b/elastica/memory_block/memory_block_rod.py @@ -1,6 +1,6 @@ __doc__ = """Create block-structure class for collection of Cosserat rod systems.""" import numpy as np -from typing import Sequence, Literal, Callable +from typing import Sequence, Literal, Callable, Type from elastica.memory_block.memory_block_rod_base import ( MemoryBlockRodBase, make_block_memory_metadata, @@ -31,7 +31,7 @@ class MemoryBlockCosseratRod( TODO: need more documentation! """ - def __init__(self, systems: Sequence, system_idx_list): + def __init__(self, systems: Sequence, system_idx_list: tuple[int]) -> None: # separate straight and ring rods system_straight_rod = [] @@ -200,7 +200,7 @@ def __init__(self, systems: Sequence, system_idx_list): # Initialize the mixin class for symplectic time-stepper. _RodSymplecticStepperMixin.__init__(self) - def _allocate_block_variables_in_nodes(self, systems: Sequence): + def _allocate_block_variables_in_nodes(self, systems: Sequence) -> None: """ This function takes system collection and allocates the variables on node for block-structure and references allocated variables back to the @@ -250,7 +250,7 @@ def _allocate_block_variables_in_nodes(self, systems: Sequence): value_type="vector", ) - def _allocate_block_variables_in_elements(self, systems: Sequence): + def _allocate_block_variables_in_elements(self, systems: Sequence) -> None: """ This function takes system collection and allocates the variables on elements for block-structure and references allocated variables back to the @@ -341,7 +341,7 @@ def _allocate_block_variables_in_elements(self, systems: Sequence): value_type="tensor", ) - def _allocate_blocks_variables_in_voronoi(self, systems: Sequence): + def _allocate_blocks_variables_in_voronoi(self, systems: Sequence) -> None: """ This function takes system collection and allocates the variables on voronoi for block-structure and references allocated variables back to the @@ -408,7 +408,9 @@ def _allocate_blocks_variables_in_voronoi(self, systems: Sequence): value_type="tensor", ) - def _allocate_blocks_variables_for_symplectic_stepper(self, systems: Sequence): + def _allocate_blocks_variables_for_symplectic_stepper( + self, systems: Sequence + ) -> None: """ This function takes system collection and allocates the variables used by symplectic stepper for block-structure and references allocated variables back to the systems. @@ -509,8 +511,8 @@ def _map_system_properties_to_block_memory( end_idx_list: np.ndarray periodic_boundary_idx: np.ndarray synchronize_periodic_boundary: Callable - domain_num: np.int64 - view_shape: tuple + domain_num: int + view_shape: tuple[int, ...] if domain_type == "node": start_idx_list = self.start_idx_in_rod_nodes.view() diff --git a/elastica/memory_block/memory_block_rod_base.py b/elastica/memory_block/memory_block_rod_base.py index c9bcc77e1..ff090c539 100644 --- a/elastica/memory_block/memory_block_rod_base.py +++ b/elastica/memory_block/memory_block_rod_base.py @@ -1,16 +1,17 @@ __doc__ = """Create block-structure class for collection of Cosserat rod systems.""" import numpy as np from typing import Iterable +import numpy.typing as npt -def make_block_memory_metadata(n_elems_in_rods: np.ndarray) -> Iterable: +def make_block_memory_metadata(n_elems_in_rods: npt.NDArray[np.integer]) -> Iterable: """ This function, takes number of elements of each rod as a numpy array and computes, ghost nodes, elements and voronoi element indexes and numbers and returns it. Parameters ---------- - n_elems_in_rods: ndarray + n_elems_in_rods: npt.NDArray An integer array containing the number of elements in each of the n rod. Returns @@ -20,9 +21,9 @@ def make_block_memory_metadata(n_elems_in_rods: np.ndarray) -> Iterable: between each pair of two rods adjacent in memory block. ghost_nodes_idx: ndarray An integer array of length n - 1 containing the indices of ghost nodes in memory block. - ghost_elements_idx: ndarray + ghost_elements_idx: npt.NDArray An integer array of length 2 * (n - 1) containing the indices of ghost elements in memory block. - ghost_voronoi_idx: ndarray + ghost_voronoi_idx: npt.NDArray An integer array of length 2 * (n - 1) containing the indices of ghost Voronoi nodes in memory block. """ @@ -52,14 +53,16 @@ def make_block_memory_metadata(n_elems_in_rods: np.ndarray) -> Iterable: return n_elems_with_ghosts, ghost_nodes_idx, ghost_elems_idx, ghost_voronoi_idx -def make_block_memory_periodic_boundary_metadata(n_elems_in_rods): +def make_block_memory_periodic_boundary_metadata( + n_elems_in_rods: npt.NDArray[np.integer], +) -> Iterable: """ This function, takes the number of elements of ring rods and computes the periodic boundary node, element and voronoi index. Parameters ---------- - n_elems_in_rods : numpy.ndarray + n_elems_in_rods : npt.NDArray 1D (n_ring_rods,) array containing data with 'float' type. Elements of this array contains total number of elements of one rod, including periodic boundary elements. @@ -67,21 +70,21 @@ def make_block_memory_periodic_boundary_metadata(n_elems_in_rods): ------- n_elems - periodic_boundary_node : numpy.ndarray + periodic_boundary_node : npt.NDArray 2D (2, n_periodic_boundary_nodes) array containing data with 'float' type. Vector containing periodic boundary elements index. First dimension is the periodic boundary index, second dimension is the referenced cell index. - periodic_boundary_elems_idx : numpy.ndarray + periodic_boundary_elems_idx : npt.NDArray 2D (2, n_periodic_boundary_elems) array containing data with 'float' type. Vector containing periodic boundary nodes index. First dimension is the periodic boundary index, second dimension is the referenced cell index. - periodic_boundary_voronoi_idx : numpy.ndarray + periodic_boundary_voronoi_idx : npt.NDArray 2D (2, n_periodic_boundary_voronoi) array containing data with 'float' type. Vector containing periodic boundary voronoi index. First dimension is the periodic boundary index, second dimension is the referenced cell index. """ - n_elem = n_elems_in_rods.copy() + n_elem: npt.NDArray[np.integer] = n_elems_in_rods.copy() n_rods = n_elems_in_rods.shape[0] periodic_boundary_node_idx = np.zeros((2, 3 * n_rods), dtype=np.int64) @@ -145,7 +148,7 @@ def make_block_memory_periodic_boundary_metadata(n_elems_in_rods): ) # Increase the n_elem in rods by 2 because we are adding two periodic boundary elements - n_elem += 2 + n_elem = n_elem + 2 return ( n_elem, @@ -160,5 +163,5 @@ class MemoryBlockRodBase: This is the base class for memory blocks for rods. """ - def __init__(self): + def __init__(self) -> None: pass From f565be140a7c677e106dcedd3e54ebf9a00be20f Mon Sep 17 00:00:00 2001 From: armantekinalp Date: Mon, 13 May 2024 14:39:29 -0500 Subject: [PATCH 043/134] fix:typing for rod and rigid body data structures to make consistent --- elastica/rigidbody/data_structures.py | 4 ++-- elastica/rod/cosserat_rod.py | 2 +- elastica/rod/data_structures.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/elastica/rigidbody/data_structures.py b/elastica/rigidbody/data_structures.py index e246c1fc8..e68f7c295 100644 --- a/elastica/rigidbody/data_structures.py +++ b/elastica/rigidbody/data_structures.py @@ -44,14 +44,14 @@ def __call__(self, time, *args, **kwargs): class _RigidRodSymplecticStepperMixin(_RodSymplecticStepperMixin): - def __init__(self): + def __init__(self) -> None: super(_RigidRodSymplecticStepperMixin, self).__init__() # Expose rate returning functions in the interface # to be used by the time-stepping algorithm # dynamic rates needs to call update_accelerations and henc # is another function - def update_internal_forces_and_torques(self, *args, **kwargs): + def update_internal_forces_and_torques(self, *args, **kwargs) -> None: pass diff --git a/elastica/rod/cosserat_rod.py b/elastica/rod/cosserat_rod.py index a6c626389..9afd5a5cd 100644 --- a/elastica/rod/cosserat_rod.py +++ b/elastica/rod/cosserat_rod.py @@ -31,7 +31,7 @@ def _get_z_vector(): return np.array([0.0, 0.0, 1.0]).reshape(3, -1) -def _compute_sigma_kappa_for_blockstructure(memory_block): +def _compute_sigma_kappa_for_blockstructure(memory_block) -> None: """ This function is a wrapper to call functions which computes shear stretch, strain and bending twist and strain. diff --git a/elastica/rod/data_structures.py b/elastica/rod/data_structures.py index 3c8b40328..2bdc3648e 100644 --- a/elastica/rod/data_structures.py +++ b/elastica/rod/data_structures.py @@ -43,7 +43,7 @@ class _RodSymplecticStepperMixin: - def __init__(self): + def __init__(self) -> None: self.kinematic_states = _KinematicState( self.position_collection, self.director_collection ) From 00b1fed3e3f54041192033510bc01b967c262641 Mon Sep 17 00:00:00 2001 From: Arman Tekinalp <53585636+armantekinalp@users.noreply.github.com> Date: Tue, 14 May 2024 19:56:39 -0500 Subject: [PATCH 044/134] Update elastica/memory_block/memory_block_rigid_body.py Co-authored-by: Seung Hyun Kim --- elastica/memory_block/memory_block_rigid_body.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/elastica/memory_block/memory_block_rigid_body.py b/elastica/memory_block/memory_block_rigid_body.py index 7c38fcda0..078cc37cc 100644 --- a/elastica/memory_block/memory_block_rigid_body.py +++ b/elastica/memory_block/memory_block_rigid_body.py @@ -7,7 +7,7 @@ class MemoryBlockRigidBody(RigidBodyBase, _RigidRodSymplecticStepperMixin): - def __init__(self, systems: Sequence, system_idx_list: Sequence[np.int64]) -> None: + def __init__(self, systems: list[RigidBodyBase], system_idx_list: list[SystemIdxType]) -> None: self.n_bodies = len(systems) self.n_elems = self.n_bodies From fdd2db9e29df840baf689ae89d2b5a94caf192df Mon Sep 17 00:00:00 2001 From: Arman Tekinalp <53585636+armantekinalp@users.noreply.github.com> Date: Tue, 14 May 2024 19:57:20 -0500 Subject: [PATCH 045/134] Update elastica/memory_block/memory_block_rod.py Co-authored-by: Seung Hyun Kim --- elastica/memory_block/memory_block_rod.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/elastica/memory_block/memory_block_rod.py b/elastica/memory_block/memory_block_rod.py index ad2ed0bbb..bee7f01ad 100644 --- a/elastica/memory_block/memory_block_rod.py +++ b/elastica/memory_block/memory_block_rod.py @@ -31,7 +31,7 @@ class MemoryBlockCosseratRod( TODO: need more documentation! """ - def __init__(self, systems: Sequence, system_idx_list: tuple[int]) -> None: + def __init__(self, systems: list[RodBase], system_idx_list: list[SystemIdxType]) -> None: # separate straight and ring rods system_straight_rod = [] From 2a7b64b14e0d2534c853959e4c0027a733854560 Mon Sep 17 00:00:00 2001 From: armantekinalp Date: Tue, 14 May 2024 20:25:26 -0500 Subject: [PATCH 046/134] update type hinting in memory block rigid body and rod and add system types --- elastica/memory_block/memory_block_rigid_body.py | 16 +++++++++------- elastica/memory_block/memory_block_rod.py | 16 +++++++++------- 2 files changed, 18 insertions(+), 14 deletions(-) diff --git a/elastica/memory_block/memory_block_rigid_body.py b/elastica/memory_block/memory_block_rigid_body.py index 078cc37cc..9ea229dc4 100644 --- a/elastica/memory_block/memory_block_rigid_body.py +++ b/elastica/memory_block/memory_block_rigid_body.py @@ -1,13 +1,15 @@ __doc__ = """Create block-structure class for collection of rigid body systems.""" import numpy as np -from typing import Sequence, Literal +from typing import Literal from elastica.rigidbody import RigidBodyBase from elastica.rigidbody.data_structures import _RigidRodSymplecticStepperMixin class MemoryBlockRigidBody(RigidBodyBase, _RigidRodSymplecticStepperMixin): - def __init__(self, systems: list[RigidBodyBase], system_idx_list: list[SystemIdxType]) -> None: + def __init__( + self, systems: list[RigidBodyBase], system_idx_list: list[SystemIdxType] + ) -> None: self.n_bodies = len(systems) self.n_elems = self.n_bodies @@ -23,7 +25,7 @@ def __init__(self, systems: list[RigidBodyBase], system_idx_list: list[SystemIdx # Initialize the mixin class for symplectic time-stepper. _RigidRodSymplecticStepperMixin.__init__(self) - def _allocate_block_variables_scalars(self, systems: Sequence) -> None: + def _allocate_block_variables_scalars(self, systems: list[RigidBodyBase]) -> None: """ This function takes system collection and allocates the variables for block-structure and references allocated variables back to the systems. @@ -58,7 +60,7 @@ def _allocate_block_variables_scalars(self, systems: Sequence) -> None: value_type="scalar", ) - def _allocate_block_variables_vectors(self, systems: Sequence) -> None: + def _allocate_block_variables_vectors(self, systems: list[RigidBodyBase]) -> None: """ This function takes system collection and allocates the vector variables for block-structure and references allocated vector variables back to the systems. @@ -94,7 +96,7 @@ def _allocate_block_variables_vectors(self, systems: Sequence) -> None: value_type="vector", ) - def _allocate_block_variables_matrix(self, systems: Sequence) -> None: + def _allocate_block_variables_matrix(self, systems: list[RigidBodyBase]) -> None: """ This function takes system collection and allocates the matrix variables for block-structure and references allocated matrix variables back to the systems. @@ -131,7 +133,7 @@ def _allocate_block_variables_matrix(self, systems: Sequence) -> None: ) def _allocate_block_variables_for_symplectic_stepper( - self, systems: Sequence + self, systems: list[RigidBodyBase] ) -> None: """ This function takes system collection and allocates the variables used by symplectic @@ -178,7 +180,7 @@ def _allocate_block_variables_for_symplectic_stepper( def _map_system_properties_to_block_memory( self, mapping_dict: dict, - systems: Sequence, + systems: list[RigidBodyBase], block_memory: np.ndarray, value_type: Literal["scalar", "vector", "tensor"], ) -> None: diff --git a/elastica/memory_block/memory_block_rod.py b/elastica/memory_block/memory_block_rod.py index bee7f01ad..e13f991ae 100644 --- a/elastica/memory_block/memory_block_rod.py +++ b/elastica/memory_block/memory_block_rod.py @@ -1,6 +1,6 @@ __doc__ = """Create block-structure class for collection of Cosserat rod systems.""" import numpy as np -from typing import Sequence, Literal, Callable, Type +from typing import Literal, Callable from elastica.memory_block.memory_block_rod_base import ( MemoryBlockRodBase, make_block_memory_metadata, @@ -31,7 +31,9 @@ class MemoryBlockCosseratRod( TODO: need more documentation! """ - def __init__(self, systems: list[RodBase], system_idx_list: list[SystemIdxType]) -> None: + def __init__( + self, systems: list[CosseratRod], system_idx_list: list[SystemIdxType] + ) -> None: # separate straight and ring rods system_straight_rod = [] @@ -200,7 +202,7 @@ def __init__(self, systems: list[RodBase], system_idx_list: list[SystemIdxType]) # Initialize the mixin class for symplectic time-stepper. _RodSymplecticStepperMixin.__init__(self) - def _allocate_block_variables_in_nodes(self, systems: Sequence) -> None: + def _allocate_block_variables_in_nodes(self, systems: list[CosseratRod]) -> None: """ This function takes system collection and allocates the variables on node for block-structure and references allocated variables back to the @@ -250,7 +252,7 @@ def _allocate_block_variables_in_nodes(self, systems: Sequence) -> None: value_type="vector", ) - def _allocate_block_variables_in_elements(self, systems: Sequence) -> None: + def _allocate_block_variables_in_elements(self, systems: list[CosseratRod]) -> None: """ This function takes system collection and allocates the variables on elements for block-structure and references allocated variables back to the @@ -341,7 +343,7 @@ def _allocate_block_variables_in_elements(self, systems: Sequence) -> None: value_type="tensor", ) - def _allocate_blocks_variables_in_voronoi(self, systems: Sequence) -> None: + def _allocate_blocks_variables_in_voronoi(self, systems: list[CosseratRod]) -> None: """ This function takes system collection and allocates the variables on voronoi for block-structure and references allocated variables back to the @@ -409,7 +411,7 @@ def _allocate_blocks_variables_in_voronoi(self, systems: Sequence) -> None: ) def _allocate_blocks_variables_for_symplectic_stepper( - self, systems: Sequence + self, systems: list[CosseratRod] ) -> None: """ This function takes system collection and allocates the variables used by symplectic @@ -477,7 +479,7 @@ def _allocate_blocks_variables_for_symplectic_stepper( def _map_system_properties_to_block_memory( self, mapping_dict: dict, - systems: Sequence, + systems: list[CosseratRod], block_memory: np.ndarray, domain_type: Literal["node", "element", "voronoi"], value_type: Literal["scalar", "vector", "tensor"], From 4d05804632fb2839dfcfad7e7c18f1eb180af7e7 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Wed, 15 May 2024 13:35:24 -0500 Subject: [PATCH 047/134] update elastica typing alias --- elastica/typing.py | 25 ++++++++++++++++++++++--- 1 file changed, 22 insertions(+), 3 deletions(-) diff --git a/elastica/typing.py b/elastica/typing.py index d8eba940f..6599f1af9 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -4,7 +4,7 @@ """ from typing import TYPE_CHECKING -from typing import Type, Union, Callable, Any +from typing import Type, Union, Callable, Any, ParamSpec from typing import TypeAlias @@ -14,7 +14,7 @@ from .rod import RodBase from .rigidbody import RigidBodyBase from .surface import SurfaceBase - from .modules import BaseSystemCollection + from .modules.base_system import BaseSystemCollection from .rod.data_structures import _State as State from .systems.protocol import SymplecticSystemProtocol, ExplicitSystemProtocol @@ -23,6 +23,14 @@ SymplecticStepperProtocol, MemoryProtocol, ) + + # Modules Base Classes + from .boundary_conditions import FreeBC + from .callback_functions import CallBackBaseClass + from .contact_forces import NoContact + from .dissipation import DamperBase + from .external_forces import NoForces + from .joint import FreeJoint else: RodBase = None RigidBodyBase = None @@ -37,8 +45,17 @@ SymplecticStepperProtocol = None MemoryProtocol = None + # Modules Base Classes + FreeBC = None + CallBackBaseClass = None + NoContact = None + DamperBase = None + NoForces = None + FreeJoint = None + SystemType: TypeAlias = Union[SymplecticSystemProtocol, ExplicitSystemProtocol] +SystemIdxType: TypeAlias = int # TODO: Modify this line and move to elastica/typing.py once system state is defined # Mostly used in explicit stepper: for symplectic, use kinetic and dynamic state @@ -83,6 +100,8 @@ SystemCollectionType: TypeAlias = BaseSystemCollection AllowedContactType: TypeAlias = Union[SystemType, Type[SurfaceBase]] +# Operators in elastica.modules +CallbackParam = ParamSpec("CallbackParam") OperatorType: TypeAlias = Callable[[float], None] -OperatorCallbackType: TypeAlias = Callable[[float, int], None] +OperatorCallbackType: TypeAlias = Callable[CallbackParam, None] OperatorFinalizeType: TypeAlias = Callable From 7506ea04f2db1e4394f273e59ec54bd9a980dd80 Mon Sep 17 00:00:00 2001 From: armantekinalp Date: Wed, 15 May 2024 13:44:07 -0500 Subject: [PATCH 048/134] typing SystemIdxType import added in memory block rigid body and rod --- elastica/memory_block/memory_block_rigid_body.py | 1 + elastica/memory_block/memory_block_rod.py | 1 + 2 files changed, 2 insertions(+) diff --git a/elastica/memory_block/memory_block_rigid_body.py b/elastica/memory_block/memory_block_rigid_body.py index 9ea229dc4..afff608e0 100644 --- a/elastica/memory_block/memory_block_rigid_body.py +++ b/elastica/memory_block/memory_block_rigid_body.py @@ -4,6 +4,7 @@ from elastica.rigidbody import RigidBodyBase from elastica.rigidbody.data_structures import _RigidRodSymplecticStepperMixin +from elastica.typing import SystemIdxType class MemoryBlockRigidBody(RigidBodyBase, _RigidRodSymplecticStepperMixin): diff --git a/elastica/memory_block/memory_block_rod.py b/elastica/memory_block/memory_block_rod.py index e13f991ae..c1b9e2381 100644 --- a/elastica/memory_block/memory_block_rod.py +++ b/elastica/memory_block/memory_block_rod.py @@ -17,6 +17,7 @@ _synchronize_periodic_boundary_of_scalar_collection, _synchronize_periodic_boundary_of_matrix_collection, ) +from elastica.typing import SystemIdxType class MemoryBlockCosseratRod( From 55f4944d515083b1c8ed71f0816582b667abe931 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Wed, 15 May 2024 13:52:27 -0500 Subject: [PATCH 049/134] fix requisite module type --- elastica/systems/protocol.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/elastica/systems/protocol.py b/elastica/systems/protocol.py index 1661a87d2..db88665ca 100644 --- a/elastica/systems/protocol.py +++ b/elastica/systems/protocol.py @@ -2,6 +2,8 @@ from typing import Protocol, Type from elastica.typing import StateType, SystemType +from elastica.modules.protocol import ModuleProtocol + from elastica.rod.data_structures import _KinematicState, _DynamicState import numpy as np @@ -13,7 +15,7 @@ class SystemProtocol(Protocol): Protocol for all elastica system """ - REQUISITE_MODULES: list[Type[SystemType]] + REQUISITE_MODULES: list[Type[ModuleProtocol]] @property def n_nodes(self) -> int: ... From 3111d79452bc7aa5ffc78b204932787a34aa879b Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Wed, 15 May 2024 23:04:22 -0500 Subject: [PATCH 050/134] add generic type for operator grouping --- elastica/modules/operator_group.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/elastica/modules/operator_group.py b/elastica/modules/operator_group.py index 96d14d900..2a3975728 100644 --- a/elastica/modules/operator_group.py +++ b/elastica/modules/operator_group.py @@ -1,4 +1,4 @@ -from typing import TypeVar +from typing import TypeVar, Generic, Iterator from collections.abc import Iterable @@ -8,7 +8,7 @@ F = TypeVar("F") -class OperatorGroupFIFO(Iterable): +class OperatorGroupFIFO(Iterable, Generic[T, F]): """ A class to store the features and their corresponding operators in a FIFO manner. @@ -40,22 +40,22 @@ class OperatorGroupFIFO(Iterable): Used to check if the specific feature is the last feature in the FIFO. """ - def __init__(self): + def __init__(self) -> None: self._operator_collection: list[list[T]] = [] self._operator_ids: list[int] = [] - def __iter__(self) -> T: + def __iter__(self) -> Iterator[T]: """Returns an operator iterator to satisfy the Iterable protocol.""" operator_chain = itertools.chain.from_iterable(self._operator_collection) for operator in operator_chain: yield operator - def append_id(self, feature: F): + def append_id(self, feature: F) -> None: """Appends the id of the feature to the list of ids.""" self._operator_ids.append(id(feature)) self._operator_collection.append([]) - def add_operators(self, feature: F, operators: list[T]): + def add_operators(self, feature: F, operators: list[T]) -> None: """Adds the operators to the list of operators corresponding to the feature.""" idx = self._operator_ids.index(id(feature)) self._operator_collection[idx].extend(operators) From 35e78f18acd2be3bc4ba1d92c3c48ab3552faf9f Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Thu, 16 May 2024 00:42:47 -0500 Subject: [PATCH 051/134] typing: connections module and necessary changes for protocol --- elastica/joint.py | 20 ++-- elastica/modules/connections.py | 166 ++++++++++++++++++-------------- elastica/modules/protocol.py | 24 +++-- elastica/typing.py | 3 +- pyproject.toml | 1 + 5 files changed, 126 insertions(+), 88 deletions(-) diff --git a/elastica/joint.py b/elastica/joint.py index e37d6f058..f4bea4019 100644 --- a/elastica/joint.py +++ b/elastica/joint.py @@ -42,18 +42,22 @@ def __init__(self, k, nu): self.nu = nu def apply_forces( - self, system_one: SystemType, index_one, system_two: SystemType, index_two + self, + system_one: SystemType, + index_one: int, + system_two: SystemType, + index_two: int, ): """ Apply joint force to the connected rod objects. Parameters ---------- - system_one : object + system_one : SystemType Rod or rigid-body object index_one : int Index of first rod for joint. - system_two : object + system_two : SystemType Rod or rigid-body object index_two : int Index of second rod for joint. @@ -81,7 +85,11 @@ def apply_forces( return def apply_torques( - self, system_one: SystemType, index_one, system_two: SystemType, index_two + self, + system_one: SystemType, + index_one: int, + system_two: SystemType, + index_two: int, ): """ Apply restoring joint torques to the connected rod objects. @@ -90,11 +98,11 @@ def apply_torques( Parameters ---------- - system_one : object + system_one : SystemType Rod or rigid-body object index_one : int Index of first rod for joint. - system_two : object + system_two : SystemType Rod or rigid-body object index_two : int Index of second rod for joint. diff --git a/elastica/modules/connections.py b/elastica/modules/connections.py index 405fe9220..6576d9688 100644 --- a/elastica/modules/connections.py +++ b/elastica/modules/connections.py @@ -5,9 +5,37 @@ Provides the connections interface to connect entities (rods, rigid bodies) using joints (see `joints.py`). """ +from typing import Type, Protocol, TypeAlias, cast +from typing_extensions import Self +from elastica.typing import ( + SystemIdxType, + OperatorParam, + OperatorFinalizeType, + SystemType, +) import numpy as np from elastica.joint import FreeJoint +from .protocol import SystemCollectionProtocol, ModuleProtocol + +ConnectionIndex: TypeAlias = ( + int | np.int_ | list[int] | tuple[int] | np.ndarray[int] | None +) + + +class SystemCollectionWithConnectionProtocol(SystemCollectionProtocol, Protocol): + _connections: list[ModuleProtocol] + + _finalize_connections: OperatorFinalizeType + + def connect( + self, + first_rod: SystemType, + second_rod: SystemType, + first_connect_idx: int = None, + second_connect_idx: int = None, + ) -> ModuleProtocol: ... + class Connections: """ @@ -21,14 +49,18 @@ class Connections: List of joint classes defined for rod-like objects. """ - def __init__(self): + def __init__(self: SystemCollectionWithConnectionProtocol): self._connections = [] super(Connections, self).__init__() self._feature_group_finalize.append(self._finalize_connections) def connect( - self, first_rod, second_rod, first_connect_idx=None, second_connect_idx=None - ): + self: SystemCollectionWithConnectionProtocol, + first_rod: SystemType, + second_rod: SystemType, + first_connect_idx: ConnectionIndex = None, + second_connect_idx: ConnectionIndex = None, + ) -> ModuleProtocol: """ This method connects two rod-like objects using the selected joint class. You need to input the two rod-like objects that are to be connected as well @@ -36,13 +68,13 @@ def connect( Parameters ---------- - first_rod : object + first_rod : SystemType Rod-like object - second_rod : object + second_rod : SystemType Rod-like object - first_connect_idx : int + first_connect_idx : Optional[int] Index of first rod for joint. - second_connect_idx : int + second_connect_idx : Optional[int] Index of second rod for joint. Returns @@ -59,14 +91,14 @@ def connect( sys_dofs = [self._systems[idx].n_elems for idx in sys_idx] # Create _Connect object, cache it and return to user - _connect = _Connect(*sys_idx, *sys_dofs) - _connect.set_index(first_connect_idx, second_connect_idx) + _connect: ModuleProtocol = _Connect(*sys_idx, *sys_dofs) + _connect.set_index(first_connect_idx, second_connect_idx) # type: ignore[attr-defined] self._connections.append(_connect) self._feature_group_synchronize.append_id(_connect) return _connect - def _finalize_connections(self): + def _finalize_connections(self: SystemCollectionWithConnectionProtocol) -> None: # From stored _Connect objects, instantiate the joints and store it # dev : the first indices stores the # (first rod index, second_rod_idx, connection_idx_on_first_rod, connection_idx_on_second_rod) @@ -76,10 +108,10 @@ def _finalize_connections(self): first_sys_idx, second_sys_idx, first_connect_idx, second_connect_idx = ( connection.id() ) - connect_instance = connection.instantiate() + connect_instance: FreeJoint = connection.instantiate() # FIXME: lambda t is included because OperatorType takes time as an argument - def apply_forces(time): + def apply_forces(time: np.floating) -> None: connect_instance.apply_forces( system_one=self._systems[first_sys_idx], index_one=first_connect_idx, @@ -87,7 +119,7 @@ def apply_forces(time): index_two=second_connect_idx, ) - def apply_torques(time): + def apply_torques(time: np.floating) -> None: connect_instance.apply_torques( system_one=self._systems[first_sys_idx], index_one=first_connect_idx, @@ -99,8 +131,6 @@ def apply_torques(time): connection, [apply_forces, apply_torques] ) - self.warnings(connection) - self._connections = [] del self._connections @@ -108,9 +138,6 @@ def apply_torques(time): # This is to optimize the call tree for better memory accesses # https://brooksandrew.github.io/simpleblog/articles/intro-to-graph-optimization-solving-cpp/ - def warnings(self, connection): - pass - class _Connect: """ @@ -118,8 +145,8 @@ class _Connect: Attributes ---------- - _first_sys_idx: int - _second_sys_idx: int + _first_sys_idx: SystemIdxType + _second_sys_idx: SystemIdxType _first_sys_n_lim: int _second_sys_n_lim: int _connect_class: list @@ -133,8 +160,8 @@ class _Connect: def __init__( self, - first_sys_idx: int, - second_sys_idx: int, + first_sys_idx: SystemIdxType, + second_sys_idx: SystemIdxType, first_sys_nlim: int, second_sys_nlim: int, ): @@ -147,71 +174,60 @@ def __init__( first_sys_nlim: int second_sys_nlim: int """ - self._first_sys_idx = first_sys_idx - self._second_sys_idx = second_sys_idx - self._first_sys_n_lim = first_sys_nlim - self._second_sys_n_lim = second_sys_nlim - self._connect_cls = None - self._args = () - self._kwargs = {} - self.first_sys_connection_idx = None - self.second_sys_connection_idx = None - - def set_index(self, first_idx, second_idx): - # TODO assert range - # First check if the types of first rod idx and second rod idx variable are same. - assert type(first_idx) is type( - second_idx - ), "Type of first_connect_idx :{}".format( - type(first_idx) - ) + " is different than second_connect_idx :{}".format( - type(second_idx) - ) + self._first_sys_idx: SystemIdxType = first_sys_idx + self._second_sys_idx: SystemIdxType = second_sys_idx + self._first_sys_n_lim: int = first_sys_nlim + self._second_sys_n_lim: int = second_sys_nlim + self._connect_cls: Type[FreeJoint] + self.first_sys_connection_idx: ConnectionIndex + self.second_sys_connection_idx: ConnectionIndex + + def set_index( + self, first_idx: ConnectionIndex, second_idx: ConnectionIndex + ) -> None: + first_type = type(first_idx) + second_type = type(second_idx) + # Check if the types of first rod idx and second rod idx variable are same. + assert ( + first_type == second_type + ), f"Type of first_connect_idx :{first_type} is different than second_connect_idx :{second_type}" # Check if the type of idx variables are correct. + allow_types = (int, np.int_, list, tuple, np.ndarray, type(None)) assert isinstance( - first_idx, (int, np.int_, list, tuple, np.ndarray, type(None)) - ), "Connection index type is not supported :{}".format( - type(first_idx) - ) + ", please try one of the following :{}".format( - (int, np.int_, list, tuple, np.ndarray) - ) + first_idx, allow_types + ), f"Connection index type is not supported :{first_type}, please try one of the following :{allow_types}" # If type of idx variables are tuple or list or np.ndarray, check validity of each entry. - if ( - isinstance(first_idx, tuple) - or isinstance(first_idx, list) - or isinstance(first_idx, np.ndarray) - ): - - for i in range(len(first_idx)): - assert isinstance(first_idx[i], (int, np.int_)), ( + if isinstance(first_idx, (tuple, list, np.ndarray)): + first_idx_ = cast(list[int], first_idx) + second_idx_ = cast(list[int], second_idx) + for i in range(len(first_idx_)): + assert isinstance(first_idx_[i], (int, np.int_)), ( "Connection index of first rod is not integer :{}".format( - first_idx[i] + first_idx_[i] ) - + " It should be :{}".format((int, np.int_)) - + " Check your input!" + + " It should be : integer. Check your input!" ) - assert isinstance(second_idx[i], (int, np.int_)), ( + assert isinstance(second_idx_[i], (int, np.int_)), ( "Connection index of second rod is not integer :{}".format( - second_idx[i] + second_idx_[i] ) - + " It should be :{}".format((int, np.int_)) - + " Check your input!" + + " It should be : integer. Check your input!" ) # The addition of +1 and and <= check on the RHS is because # connections can be made to the node indices as well assert ( -(self._first_sys_n_lim + 1) - <= first_idx[i] + <= first_idx_[i] <= self._first_sys_n_lim ), "Connection index of first rod exceeds its dof : {}".format( self._first_sys_n_lim ) assert ( -(self._second_sys_n_lim + 1) - <= second_idx[i] + <= second_idx_[i] <= self._second_sys_n_lim ), "Connection index of second rod exceeds its dof : {}".format( self._second_sys_n_lim @@ -220,16 +236,17 @@ def set_index(self, first_idx, second_idx): # Do nothing if idx are None pass else: - # The addition of +1 and and <= check on the RHS is because # connections can be made to the node indices as well + first_idx__ = cast(int, first_idx) + second_idx__ = cast(int, second_idx) assert ( - -(self._first_sys_n_lim + 1) <= first_idx <= self._first_sys_n_lim + -(self._first_sys_n_lim + 1) <= first_idx__ <= self._first_sys_n_lim ), "Connection index of first rod exceeds its dof : {}".format( self._first_sys_n_lim ) assert ( - -(self._second_sys_n_lim + 1) <= second_idx <= self._second_sys_n_lim + -(self._second_sys_n_lim + 1) <= second_idx__ <= self._second_sys_n_lim ), "Connection index of second rod exceeds its dof : {}".format( self._second_sys_n_lim ) @@ -237,7 +254,12 @@ def set_index(self, first_idx, second_idx): self.first_sys_connection_idx = first_idx self.second_sys_connection_idx = second_idx - def using(self, connect_cls, *args, **kwargs): + def using( + self, + connect_cls: Type[FreeJoint], + *args: OperatorParam.args, + **kwargs: OperatorParam.kwargs, + ) -> Self: """ This method is a module to set which joint class is used to connect user defined rod-like objects. @@ -265,7 +287,9 @@ def using(self, connect_cls, *args, **kwargs): self._kwargs = kwargs return self - def id(self): + def id( + self, + ) -> tuple[SystemIdxType, SystemIdxType, ConnectionIndex, ConnectionIndex]: return ( self._first_sys_idx, self._second_sys_idx, @@ -273,7 +297,7 @@ def id(self): self.second_sys_connection_idx, ) - def instantiate(self): + def instantiate(self) -> FreeJoint: if not self._connect_cls: raise RuntimeError( "No connections provided to link rod id {0}" diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index 5eaf7d85f..424423d0a 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -1,6 +1,5 @@ -from typing import Protocol, Type, Generator, Iterable +from typing import Protocol, Generator from typing_extensions import Self # 3.11: from typing import Self - from elastica.typing import ( SystemIdxType, OperatorType, @@ -9,33 +8,38 @@ SystemType, ) +from .operator_group import OperatorGroupFIFO + class SystemCollectionProtocol(Protocol): _systems: list[SystemType] @property - def _feature_group_synchronize(self) -> Iterable[OperatorType]: ... + def _feature_group_synchronize(self) -> OperatorGroupFIFO: ... @property - def _feature_group_constrain_values(self) -> Iterable[OperatorType]: ... + def _feature_group_constrain_values(self) -> list[OperatorType]: ... @property - def _feature_group_constrain_rates(self) -> Iterable[OperatorType]: ... + def _feature_group_constrain_rates(self) -> list[OperatorType]: ... @property - def _feature_group_callback(self) -> Iterable[OperatorCallbackType]: ... + def _feature_group_callback(self) -> list[OperatorCallbackType]: ... @property - def _feature_group_finalize(self) -> Iterable[OperatorFinalizeType]: ... + def _feature_group_finalize(self) -> list[OperatorFinalizeType]: ... def blocks(self) -> Generator[SystemType, None, None]: ... def _get_sys_idx_if_valid(self, sys_to_be_added: SystemType) -> SystemIdxType: ... -class ModuleProtocol(Protocol): +M = TypeVar("M", bound=ModuleProtocol) + + +class ModuleProtocol(Protocol[M]): def using(self, callback_cls, *args, **kwargs) -> Self: ... - def instantiate(self) -> Type: ... + def instantiate(self) -> M: ... - def id(self) -> SystemIdxType: ... + def id(self) -> Any: ... diff --git a/elastica/typing.py b/elastica/typing.py index 3f4b0773b..b20d3c999 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -107,7 +107,8 @@ AllowedContactType: TypeAlias = Union[SystemType, Type[SurfaceBase]] # Operators in elastica.modules +OperatorParam = ParamSpec("OperatorParam") CallbackParam = ParamSpec("CallbackParam") -OperatorType: TypeAlias = Callable[[float], None] +OperatorType: TypeAlias = Callable[OperatorParam, None] OperatorCallbackType: TypeAlias = Callable[CallbackParam, None] OperatorFinalizeType: TypeAlias = Callable diff --git a/pyproject.toml b/pyproject.toml index 09a13f1ec..0816b9caa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -122,6 +122,7 @@ strict = true allow_redefinition = false check_untyped_defs = true +disallow_any_unimported = true disallow_any_generics = false disallow_incomplete_defs = true disallow_untyped_calls = true From 1418d67079c3e0125b377b2b72cb00259a5f6a26 Mon Sep 17 00:00:00 2001 From: Ankith Date: Sun, 19 May 2024 22:20:10 +0530 Subject: [PATCH 052/134] Improve typehinting at root directory --- elastica/_calculus.py | 32 ++- elastica/_contact_functions.py | 213 ++++++++++---------- elastica/_linalg.py | 45 +++-- elastica/_rotations.py | 48 +++-- elastica/_synchronize_periodic_boundary.py | 23 ++- elastica/boundary_conditions.py | 113 +++++++---- elastica/callback_functions.py | 31 ++- elastica/contact_forces.py | 58 +++--- elastica/contact_utils.py | 89 ++++---- elastica/dissipation.py | 47 +++-- elastica/external_forces.py | 125 +++++++----- elastica/interaction.py | 172 +++++++++------- elastica/joint.py | 223 ++++++++++++--------- elastica/restart.py | 21 +- elastica/transformations.py | 21 +- elastica/typing.py | 2 +- elastica/utils.py | 28 ++- 17 files changed, 765 insertions(+), 526 deletions(-) diff --git a/elastica/_calculus.py b/elastica/_calculus.py index eca9829b7..ae730da15 100644 --- a/elastica/_calculus.py +++ b/elastica/_calculus.py @@ -1,6 +1,8 @@ __doc__ = """ Quadrature and difference kernels """ +from typing import Any, Union import numpy as np from numpy import zeros, empty +from numpy.typing import NDArray from numba import njit from elastica.reset_functions_for_block_structure._reset_ghost_vector_or_scalar import ( _reset_vector_ghost, @@ -9,15 +11,17 @@ @functools.lru_cache(maxsize=2) -def _get_zero_array(dim, ndim): +def _get_zero_array(dim: int, ndim: int) -> Union[float, NDArray[np.floating], None]: if ndim == 1: return 0.0 if ndim == 2: return np.zeros((dim, 1)) + return None + @njit(cache=True) -def _trapezoidal(array_collection): +def _trapezoidal(array_collection: NDArray[np.floating]) -> NDArray[np.floating]: """ Simple trapezoidal quadrature rule with zero at end-points, in a dimension agnostic way @@ -63,7 +67,9 @@ def _trapezoidal(array_collection): @njit(cache=True) -def _trapezoidal_for_block_structure(array_collection, ghost_idx): +def _trapezoidal_for_block_structure( + array_collection: NDArray[np.floating], ghost_idx: NDArray[np.integer] +) -> NDArray[np.floating]: """ Simple trapezoidal quadrature rule with zero at end-points, in a dimension agnostic way. This form specifically for the block structure implementation and there is a reset function call, to reset @@ -115,7 +121,9 @@ def _trapezoidal_for_block_structure(array_collection, ghost_idx): @njit(cache=True) -def _two_point_difference(array_collection): +def _two_point_difference( + array_collection: NDArray[np.floating], +) -> NDArray[np.floating]: """ This function does differentiation. @@ -156,7 +164,9 @@ def _two_point_difference(array_collection): @njit(cache=True) -def _two_point_difference_for_block_structure(array_collection, ghost_idx): +def _two_point_difference_for_block_structure( + array_collection: NDArray[np.floating], ghost_idx: NDArray[np.integer] +) -> NDArray[np.floating]: """ This function does the differentiation, for Cosserat rod model equations. This form specifically for the block structure implementation and there is a reset function call, to @@ -207,7 +217,7 @@ def _two_point_difference_for_block_structure(array_collection, ghost_idx): @njit(cache=True) -def _difference(vector): +def _difference(vector: NDArray[np.floating]) -> NDArray[np.floating]: """ This function computes difference between elements of a batch vector. @@ -238,7 +248,7 @@ def _difference(vector): @njit(cache=True) -def _average(vector): +def _average(vector: NDArray[np.floating]) -> NDArray[np.floating]: """ This function computes the average between elements of a vector. @@ -268,7 +278,9 @@ def _average(vector): @njit(cache=True) -def _clip_array(input_array, vmin, vmax): +def _clip_array( + input_array: NDArray[np.floating], vmin: np.floating, vmax: np.floating +) -> NDArray[np.floating]: """ This function clips an array values between user defined minimum and maximum @@ -304,7 +316,7 @@ def _clip_array(input_array, vmin, vmax): @njit(cache=True) -def _isnan_check(array): +def _isnan_check(array: NDArray[Any]) -> bool: """ This function checks if there is any nan inside the array. If there is nan, it returns True boolean. @@ -324,7 +336,7 @@ def _isnan_check(array): Python version: 2.24 µs ± 96.1 ns per loop This version: 479 ns ± 6.49 ns per loop """ - return np.isnan(array).any() + return bool(np.isnan(array).any()) position_difference_kernel = _difference diff --git a/elastica/_contact_functions.py b/elastica/_contact_functions.py index 245d92319..3b2d25d5a 100644 --- a/elastica/_contact_functions.py +++ b/elastica/_contact_functions.py @@ -24,28 +24,29 @@ ) import numba import numpy as np +from numpy.typing import NDArray @numba.njit(cache=True) def _calculate_contact_forces_rod_cylinder( - x_collection_rod, - edge_collection_rod, - x_cylinder_center, - x_cylinder_tip, - edge_cylinder, - radii_sum, - length_sum, - internal_forces_rod, - external_forces_rod, - external_forces_cylinder, - external_torques_cylinder, - cylinder_director_collection, - velocity_rod, - velocity_cylinder, - contact_k, - contact_nu, - velocity_damping_coefficient, - friction_coefficient, + x_collection_rod: NDArray[np.floating], + edge_collection_rod: NDArray[np.floating], + x_cylinder_center: NDArray[np.floating], + x_cylinder_tip: NDArray[np.floating], + edge_cylinder: NDArray[np.floating], + radii_sum: NDArray[np.floating], + length_sum: NDArray[np.floating], + internal_forces_rod: NDArray[np.floating], + external_forces_rod: NDArray[np.floating], + external_forces_cylinder: NDArray[np.floating], + external_torques_cylinder: NDArray[np.floating], + cylinder_director_collection: NDArray[np.floating], + velocity_rod: NDArray[np.floating], + velocity_cylinder: NDArray[np.floating], + contact_k: np.floating, + contact_nu: np.floating, + velocity_damping_coefficient: np.floating, + friction_coefficient: np.floating, ) -> None: # We already pass in only the first n_elem x n_points = x_collection_rod.shape[1] @@ -155,22 +156,22 @@ def _calculate_contact_forces_rod_cylinder( @numba.njit(cache=True) def _calculate_contact_forces_rod_rod( - x_collection_rod_one, - radius_rod_one, - length_rod_one, - tangent_rod_one, - velocity_rod_one, - internal_forces_rod_one, - external_forces_rod_one, - x_collection_rod_two, - radius_rod_two, - length_rod_two, - tangent_rod_two, - velocity_rod_two, - internal_forces_rod_two, - external_forces_rod_two, - contact_k, - contact_nu, + x_collection_rod_one: NDArray[np.floating], + radius_rod_one: NDArray[np.floating], + length_rod_one: NDArray[np.floating], + tangent_rod_one: NDArray[np.floating], + velocity_rod_one: NDArray[np.floating], + internal_forces_rod_one: NDArray[np.floating], + external_forces_rod_one: NDArray[np.floating], + x_collection_rod_two: NDArray[np.floating], + radius_rod_two: NDArray[np.floating], + length_rod_two: NDArray[np.floating], + tangent_rod_two: NDArray[np.floating], + velocity_rod_two: NDArray[np.floating], + internal_forces_rod_two: NDArray[np.floating], + external_forces_rod_two: NDArray[np.floating], + contact_k: np.floating, + contact_nu: np.floating, ) -> None: # We already pass in only the first n_elem x n_points_rod_one = x_collection_rod_one.shape[1] @@ -272,14 +273,14 @@ def _calculate_contact_forces_rod_rod( @numba.njit(cache=True) def _calculate_contact_forces_self_rod( - x_collection_rod, - radius_rod, - length_rod, - tangent_rod, - velocity_rod, - external_forces_rod, - contact_k, - contact_nu, + x_collection_rod: NDArray[np.floating], + radius_rod: NDArray[np.floating], + length_rod: NDArray[np.floating], + tangent_rod: NDArray[np.floating], + velocity_rod: NDArray[np.floating], + external_forces_rod: NDArray[np.floating], + contact_k: np.floating, + contact_nu: np.floating, ) -> None: # We already pass in only the first n_elem x n_points_rod = x_collection_rod.shape[1] @@ -360,24 +361,24 @@ def _calculate_contact_forces_self_rod( @numba.njit(cache=True) def _calculate_contact_forces_rod_sphere( - x_collection_rod, - edge_collection_rod, - x_sphere_center, - x_sphere_tip, - edge_sphere, - radii_sum, - length_sum, - internal_forces_rod, - external_forces_rod, - external_forces_sphere, - external_torques_sphere, - sphere_director_collection, - velocity_rod, - velocity_sphere, - contact_k, - contact_nu, - velocity_damping_coefficient, - friction_coefficient, + x_collection_rod: NDArray[np.floating], + edge_collection_rod: NDArray[np.floating], + x_sphere_center: NDArray[np.floating], + x_sphere_tip: NDArray[np.floating], + edge_sphere: NDArray[np.floating], + radii_sum: NDArray[np.floating], + length_sum: NDArray[np.floating], + internal_forces_rod: NDArray[np.floating], + external_forces_rod: NDArray[np.floating], + external_forces_sphere: NDArray[np.floating], + external_torques_sphere: NDArray[np.floating], + sphere_director_collection: NDArray[np.floating], + velocity_rod: NDArray[np.floating], + velocity_sphere: NDArray[np.floating], + contact_k: np.floating, + contact_nu: np.floating, + velocity_damping_coefficient: np.floating, + friction_coefficient: np.floating, ) -> None: # We already pass in only the first n_elem x n_points = x_collection_rod.shape[1] @@ -486,18 +487,18 @@ def _calculate_contact_forces_rod_sphere( @numba.njit(cache=True) def _calculate_contact_forces_rod_plane( - plane_origin, - plane_normal, - surface_tol, - k, - nu, - radius, - mass, - position_collection, - velocity_collection, - internal_forces, - external_forces, -): + plane_origin: NDArray[np.floating], + plane_normal: NDArray[np.floating], + surface_tol: np.floating, + k: np.floating, + nu: np.floating, + radius: NDArray[np.floating], + mass: NDArray[np.floating], + position_collection: NDArray[np.floating], + velocity_collection: NDArray[np.floating], + internal_forces: NDArray[np.floating], + external_forces: NDArray[np.floating], +) -> tuple[NDArray[np.floating], NDArray[np.intp]]: """ This function computes the plane force response on the element, in the case of contact. Contact model given in Eqn 4.8 Gazzola et. al. RSoS 2018 paper @@ -571,30 +572,30 @@ def _calculate_contact_forces_rod_plane( @numba.njit(cache=True) def _calculate_contact_forces_rod_plane_with_anisotropic_friction( - plane_origin, - plane_normal, - surface_tol, - slip_velocity_tol, - k, - nu, - kinetic_mu_forward, - kinetic_mu_backward, - kinetic_mu_sideways, - static_mu_forward, - static_mu_backward, - static_mu_sideways, - radius, - mass, - tangents, - position_collection, - director_collection, - velocity_collection, - omega_collection, - internal_forces, - external_forces, - internal_torques, - external_torques, -): + plane_origin: NDArray[np.floating], + plane_normal: NDArray[np.floating], + surface_tol: np.floating, + slip_velocity_tol: np.floating, + k: np.floating, + nu: np.floating, + kinetic_mu_forward: np.floating, + kinetic_mu_backward: np.floating, + kinetic_mu_sideways: np.floating, + static_mu_forward: np.floating, + static_mu_backward: np.floating, + static_mu_sideways: np.floating, + radius: NDArray[np.floating], + mass: NDArray[np.floating], + tangents: NDArray[np.floating], + position_collection: NDArray[np.floating], + director_collection: NDArray[np.floating], + velocity_collection: NDArray[np.floating], + omega_collection: NDArray[np.floating], + internal_forces: NDArray[np.floating], + external_forces: NDArray[np.floating], + internal_torques: NDArray[np.floating], + external_torques: NDArray[np.floating], +) -> None: ( plane_response_force_mag, no_contact_point_idx, @@ -784,16 +785,16 @@ def _calculate_contact_forces_rod_plane_with_anisotropic_friction( @numba.njit(cache=True) def _calculate_contact_forces_cylinder_plane( - plane_origin, - plane_normal, - surface_tol, - k, - nu, - length, - position_collection, - velocity_collection, - external_forces, -): + plane_origin: NDArray[np.floating], + plane_normal: NDArray[np.floating], + surface_tol: np.floating, + k: np.floating, + nu: np.floating, + length: NDArray[np.floating], + position_collection: NDArray[np.floating], + velocity_collection: NDArray[np.floating], + external_forces: NDArray[np.floating], +) -> tuple[NDArray[np.floating], NDArray[np.intp]]: # Compute plane response force # total_forces = system.internal_forces + system.external_forces diff --git a/elastica/_linalg.py b/elastica/_linalg.py index a4995ab25..3123ff75b 100644 --- a/elastica/_linalg.py +++ b/elastica/_linalg.py @@ -1,5 +1,6 @@ __doc__ = """ Convenient linear algebra kernels """ import numpy as np +from numpy.typing import NDArray from numba import njit from numpy import sqrt import functools @@ -8,7 +9,7 @@ @functools.lru_cache(maxsize=1) -def levi_civita_tensor(dim): +def levi_civita_tensor(dim: int) -> NDArray[np.floating]: """ Parameters @@ -28,7 +29,9 @@ def levi_civita_tensor(dim): @njit(cache=True) -def _batch_matvec(matrix_collection, vector_collection): +def _batch_matvec( + matrix_collection: NDArray[np.floating], vector_collection: NDArray[np.floating] +) -> NDArray[np.floating]: """ This function does batch matrix and batch vector product @@ -59,7 +62,10 @@ def _batch_matvec(matrix_collection, vector_collection): @njit(cache=True) -def _batch_matmul(first_matrix_collection, second_matrix_collection): +def _batch_matmul( + first_matrix_collection: NDArray[np.floating], + second_matrix_collection: NDArray[np.floating], +) -> NDArray[np.floating]: """ This is batch matrix matrix multiplication function. Only batch of 3x3 matrices can be multiplied. @@ -93,7 +99,10 @@ def _batch_matmul(first_matrix_collection, second_matrix_collection): @njit(cache=True) -def _batch_cross(first_vector_collection, second_vector_collection): +def _batch_cross( + first_vector_collection: NDArray[np.floating], + second_vector_collection: NDArray[np.floating], +) -> NDArray[np.floating]: """ This function does cross product between two batch vectors. @@ -133,7 +142,9 @@ def _batch_cross(first_vector_collection, second_vector_collection): @njit(cache=True) -def _batch_vec_oneD_vec_cross(first_vector_collection, second_vector): +def _batch_vec_oneD_vec_cross( + first_vector_collection: NDArray[np.floating], second_vector: NDArray[np.floating] +) -> NDArray[np.floating]: """ This function does cross product between batch vector and a 1D vector. Idea of having this function is that, for friction calculations, we dont @@ -177,7 +188,9 @@ def _batch_vec_oneD_vec_cross(first_vector_collection, second_vector): @njit(cache=True) -def _batch_dot(first_vector, second_vector): +def _batch_dot( + first_vector: NDArray[np.floating], second_vector: NDArray[np.floating] +) -> NDArray[np.floating]: """ This function does batch vec and batch vec dot product. Parameters @@ -204,7 +217,7 @@ def _batch_dot(first_vector, second_vector): @njit(cache=True) -def _batch_norm(vector): +def _batch_norm(vector: NDArray[np.floating]) -> NDArray[np.floating]: """ This function computes norm of a batch vector Parameters @@ -233,7 +246,9 @@ def _batch_norm(vector): @njit(cache=True) -def _batch_product_i_k_to_ik(vector1, vector2): +def _batch_product_i_k_to_ik( + vector1: NDArray[np.floating], vector2: NDArray[np.floating] +) -> NDArray[np.floating]: """ This function does outer product following 'i,k->ik'. vector1 has shape of 3 and vector 2 has shape of blocksize @@ -262,7 +277,9 @@ def _batch_product_i_k_to_ik(vector1, vector2): @njit(cache=True) -def _batch_product_i_ik_to_k(vector1, vector2): +def _batch_product_i_ik_to_k( + vector1: NDArray[np.floating], vector2: NDArray[np.floating] +) -> NDArray[np.floating]: """ This function does the following product 'i,ik->k' This function do dot product between a vector of 3 elements @@ -293,7 +310,9 @@ def _batch_product_i_ik_to_k(vector1, vector2): @njit(cache=True) -def _batch_product_k_ik_to_ik(vector1, vector2): +def _batch_product_k_ik_to_ik( + vector1: NDArray[np.floating], vector2: NDArray[np.floating] +) -> NDArray[np.floating]: """ This function does the following product 'k, ik->ik' Parameters @@ -322,7 +341,9 @@ def _batch_product_k_ik_to_ik(vector1, vector2): @njit(cache=True) -def _batch_vector_sum(vector1, vector2): +def _batch_vector_sum( + vector1: NDArray[np.floating], vector2: NDArray[np.floating] +) -> NDArray[np.floating]: """ This function is for summing up two vectors. Although this function is not faster than pure python implementation @@ -352,7 +373,7 @@ def _batch_vector_sum(vector1, vector2): @njit(cache=True) -def _batch_matrix_transpose(input_matrix): +def _batch_matrix_transpose(input_matrix: NDArray[np.floating]) -> NDArray[np.floating]: """ This function takes an batch input matrix and transpose it. Parameters diff --git a/elastica/_rotations.py b/elastica/_rotations.py index 25ec14212..7fb9d0633 100644 --- a/elastica/_rotations.py +++ b/elastica/_rotations.py @@ -8,6 +8,7 @@ from numpy import cos from numpy import sqrt from numpy import arccos +from numpy.typing import NDArray from numba import njit @@ -15,7 +16,9 @@ @njit(cache=True) -def _get_rotation_matrix(scale: float, axis_collection): +def _get_rotation_matrix( + scale: np.floating, axis_collection: NDArray[np.floating] +) -> NDArray[np.floating]: blocksize = axis_collection.shape[1] rot_mat = np.empty((3, 3, blocksize)) @@ -49,7 +52,11 @@ def _get_rotation_matrix(scale: float, axis_collection): @njit(cache=True) -def _rotate(director_collection, scale: float, axis_collection): +def _rotate( + director_collection: NDArray[np.floating], + scale: np.floating, + axis_collection: NDArray[np.floating], +) -> NDArray[np.floating]: """ Does alibi rotations https://en.wikipedia.org/wiki/Rotation_matrix#Ambiguities @@ -74,7 +81,7 @@ def _rotate(director_collection, scale: float, axis_collection): @njit(cache=True) -def _inv_rotate(director_collection): +def _inv_rotate(director_collection: NDArray[np.floating]) -> NDArray[np.floating]: """ Calculated rate of change using Rodrigues' formula @@ -156,12 +163,15 @@ def _inv_rotate(director_collection): return vector_collection +_generate_skew_map_sentinel = (0, 0, 0) + + # TODO: Below contains numpy-only implementations @functools.lru_cache(maxsize=1) -def _generate_skew_map(dim: int): +def _generate_skew_map(dim: int) -> list[tuple[int, int, int]]: # TODO Documentation # Preallocate - mapping_list = [None] * ((dim**2 - dim) // 2) + mapping_list = [_generate_skew_map_sentinel] * ((dim**2 - dim) // 2) # Indexing (i,j), j is the fastest changing # r = 2, r here is rank, we deal with only matrices for index, (i, j) in enumerate(combinations(range(dim), r=2)): @@ -185,7 +195,7 @@ def _generate_skew_map(dim: int): @functools.lru_cache(maxsize=1) -def _get_skew_map(dim): +def _get_skew_map(dim: int) -> tuple[tuple[int, int, int], ...]: """Generates mapping from src to target skew-symmetric operator For input vector V and output Matrix M (represented in lexicographical index), @@ -208,7 +218,7 @@ def _get_skew_map(dim): @functools.lru_cache(maxsize=1) -def _get_inv_skew_map(dim): +def _get_inv_skew_map(dim: int) -> tuple[tuple[int, int, int], ...]: # TODO Documentation # (vec_src, mat_i, mat_j, sign) mapping_list = _generate_skew_map(dim) @@ -219,7 +229,7 @@ def _get_inv_skew_map(dim): @functools.lru_cache(maxsize=1) -def _get_diag_map(dim): +def _get_diag_map(dim: int) -> tuple[int, ...]: """Generates lexicographic mapping to diagonal in a serialized matrix-type For input dimension dim we calculate mapping to * in Matrix M below @@ -231,17 +241,10 @@ def _get_diag_map(dim): in a dimension agnostic way. """ - # Preallocate - mapping_list = [None] * dim - - # Store linear indices - for dim_iter in range(dim): - mapping_list[dim_iter] = dim_iter * (dim + 1) - - return tuple(mapping_list) + return tuple([dim_iter * (dim + 1) for dim_iter in range(dim)]) -def _skew_symmetrize(vector): +def _skew_symmetrize(vector: NDArray[np.floating]) -> NDArray[np.floating]: """ Parameters @@ -276,7 +279,7 @@ def _skew_symmetrize(vector): # This is purely for testing and optimization sake # While calculating u^2, use u with einsum instead, as it is tad bit faster -def _skew_symmetrize_sq(vector): +def _skew_symmetrize_sq(vector: NDArray[np.floating]) -> NDArray[np.floating]: """ Generate the square of an orthogonal matrix from vector elements @@ -298,12 +301,11 @@ def _skew_symmetrize_sq(vector): hardcoded : 23.1 µs ± 481 ns per loop this version: 14.1 µs ± 96.9 ns per loop """ - dim, _ = vector.shape # First generate array of [x^2, xy, xz, yx, y^2, yz, zx, zy, z^2] # across blocksize # This is slightly faster than doing v[np.newaxis,:,:] * v[:,np.newaxis,:] - products_xy = np.einsum("ik,jk->ijk", vector, vector) + products_xy: NDArray[np.floating] = np.einsum("ik,jk->ijk", vector, vector) # No copy made here, as we do not change memory layout # products_xy = products_xy.reshape((dim * dim, -1)) @@ -335,7 +337,9 @@ def _skew_symmetrize_sq(vector): return products_xy -def _get_skew_symmetric_pair(vector_collection): +def _get_skew_symmetric_pair( + vector_collection: NDArray[np.floating], +) -> tuple[NDArray[np.floating], NDArray[np.floating]]: """ Parameters @@ -351,7 +355,7 @@ def _get_skew_symmetric_pair(vector_collection): return u, u_sq -def _inv_skew_symmetrize(matrix): +def _inv_skew_symmetrize(matrix: NDArray[np.floating]) -> NDArray[np.floating]: """ Return the vector elements from a skew-symmetric matrix M diff --git a/elastica/_synchronize_periodic_boundary.py b/elastica/_synchronize_periodic_boundary.py index b4fe87b4f..0a06622c8 100644 --- a/elastica/_synchronize_periodic_boundary.py +++ b/elastica/_synchronize_periodic_boundary.py @@ -2,12 +2,18 @@ """These functions are used to synchronize periodic boundaries for ring rods. """ ) +from typing import Any from numba import njit +import numpy as np +from numpy.typing import NDArray from elastica.boundary_conditions import ConstraintBase +from elastica.typing import SystemType @njit(cache=True) -def _synchronize_periodic_boundary_of_vector_collection(input, periodic_idx): +def _synchronize_periodic_boundary_of_vector_collection( + input: NDArray[np.floating], periodic_idx: NDArray[np.floating] +) -> None: """ This function synchronizes the periodic boundaries of a vector collection. Parameters @@ -28,7 +34,9 @@ def _synchronize_periodic_boundary_of_vector_collection(input, periodic_idx): @njit(cache=True) -def _synchronize_periodic_boundary_of_matrix_collection(input, periodic_idx): +def _synchronize_periodic_boundary_of_matrix_collection( + input: NDArray[np.floating], periodic_idx: NDArray[np.floating] +) -> None: """ This function synchronizes the periodic boundaries of a matrix collection. Parameters @@ -50,7 +58,9 @@ def _synchronize_periodic_boundary_of_matrix_collection(input, periodic_idx): @njit(cache=True) -def _synchronize_periodic_boundary_of_scalar_collection(input, periodic_idx): +def _synchronize_periodic_boundary_of_scalar_collection( + input: NDArray[np.floating], periodic_idx: NDArray[np.floating] +) -> None: """ This function synchronizes the periodic boundaries of a scalar collection. @@ -76,10 +86,11 @@ class _ConstrainPeriodicBoundaries(ConstraintBase): is to synchronize periodic boundaries of ring rod. """ - def __init__(self, **kwargs): + # TODO: improve typing + def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) - def constrain_values(self, rod, time): + def constrain_values(self, rod: SystemType, time: np.floating) -> None: _synchronize_periodic_boundary_of_vector_collection( rod.position_collection, rod.periodic_boundary_nodes_idx ) @@ -87,7 +98,7 @@ def constrain_values(self, rod, time): rod.director_collection, rod.periodic_boundary_elems_idx ) - def constrain_rates(self, rod, time): + def constrain_rates(self, rod: SystemType, time: np.floating) -> None: _synchronize_periodic_boundary_of_vector_collection( rod.velocity_collection, rod.periodic_boundary_nodes_idx ) diff --git a/elastica/boundary_conditions.py b/elastica/boundary_conditions.py index 0cf666c8c..7f87258ea 100644 --- a/elastica/boundary_conditions.py +++ b/elastica/boundary_conditions.py @@ -1,9 +1,10 @@ __doc__ = """ Built-in boundary condition implementationss """ import warnings -from typing import Optional +from typing import Any, Optional import numpy as np +from numpy.typing import NDArray from abc import ABC, abstractmethod @@ -34,7 +35,7 @@ class ConstraintBase(ABC): _constrained_position_idx: np.ndarray _constrained_director_idx: np.ndarray - def __init__(self, *args, **kwargs): + def __init__(self, *args: Any, **kwargs: Any) -> None: """Initialize boundary condition""" try: self._system = kwargs["_system"] @@ -67,7 +68,7 @@ def constrained_director_idx(self) -> Optional[np.ndarray]: return self._constrained_director_idx @abstractmethod - def constrain_values(self, system: SystemType, time: float) -> None: + def constrain_values(self, system: SystemType, time: np.floating) -> None: # TODO: In the future, we can remove rod and use self.system """ Constrain values (position and/or directors) of a rod object. @@ -82,7 +83,7 @@ def constrain_values(self, system: SystemType, time: float) -> None: pass @abstractmethod - def constrain_rates(self, system: SystemType, time: float) -> None: + def constrain_rates(self, system: SystemType, time: np.floating) -> None: # TODO: In the future, we can remove rod and use self.system """ Constrain rates (velocity and/or omega) of a rod object. @@ -103,14 +104,14 @@ class FreeBC(ConstraintBase): Boundary condition template. """ - def __init__(self, **kwargs): + def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) - def constrain_values(self, system: SystemType, time: float) -> None: + def constrain_values(self, system: SystemType, time: np.floating) -> None: """In FreeBC, this routine simply passes.""" pass - def constrain_rates(self, system: SystemType, time: float) -> None: + def constrain_rates(self, system: SystemType, time: np.floating) -> None: """In FreeBC, this routine simply passes.""" pass @@ -143,7 +144,12 @@ class OneEndFixedBC(ConstraintBase): ... ) """ - def __init__(self, fixed_position, fixed_directors, **kwargs): + def __init__( + self, + fixed_position: tuple[int, ...], + fixed_directors: tuple[int, ...], + **kwargs: Any, + ) -> None: """ Initialization of the constraint. Any parameter passed to 'using' will be available in kwargs. @@ -159,7 +165,7 @@ def __init__(self, fixed_position, fixed_directors, **kwargs): self.fixed_position_collection = np.array(fixed_position) self.fixed_directors_collection = np.array(fixed_directors) - def constrain_values(self, system: SystemType, time: float) -> None: + def constrain_values(self, system: SystemType, time: np.floating) -> None: # system.position_collection[..., 0] = self.fixed_position # system.director_collection[..., 0] = self.fixed_directors self.compute_constrain_values( @@ -169,7 +175,7 @@ def constrain_values(self, system: SystemType, time: float) -> None: self.fixed_directors_collection, ) - def constrain_rates(self, system: SystemType, time: float) -> None: + def constrain_rates(self, system: SystemType, time: np.floating) -> None: # system.velocity_collection[..., 0] = 0.0 # system.omega_collection[..., 0] = 0.0 self.compute_constrain_rates( @@ -180,11 +186,11 @@ def constrain_rates(self, system: SystemType, time: float) -> None: @staticmethod @njit(cache=True) def compute_constrain_values( - position_collection, - fixed_position_collection, - director_collection, - fixed_directors_collection, - ): + position_collection: NDArray[np.floating], + fixed_position_collection: NDArray[np.floating], + director_collection: NDArray[np.floating], + fixed_directors_collection: NDArray[np.floating], + ) -> None: """ Computes constrain values in numba njit decorator @@ -208,7 +214,10 @@ def compute_constrain_values( @staticmethod @njit(cache=True) - def compute_constrain_rates(velocity_collection, omega_collection): + def compute_constrain_rates( + velocity_collection: NDArray[np.floating], + omega_collection: NDArray[np.floating], + ) -> None: """ Compute contrain rates in numba njit decorator @@ -266,11 +275,11 @@ class GeneralConstraint(ConstraintBase): def __init__( self, - *fixed_data, - translational_constraint_selector: Optional[np.ndarray] = None, - rotational_constraint_selector: Optional[np.array] = None, - **kwargs, - ): + *fixed_data: Any, + translational_constraint_selector: Optional[NDArray[np.bool_]] = None, + rotational_constraint_selector: Optional[NDArray[np.bool_]] = None, + **kwargs: Any, + ) -> None: """ Initialization of the constraint. Any parameter passed to 'using' will be available in kwargs. @@ -331,7 +340,7 @@ def __init__( ) self.rotational_constraint_selector = rotational_constraint_selector.astype(int) - def constrain_values(self, system: SystemType, time: float) -> None: + def constrain_values(self, system: SystemType, time: np.floating) -> None: if self.constrained_position_idx.size: self.nb_constrain_translational_values( system.position_collection, @@ -340,7 +349,7 @@ def constrain_values(self, system: SystemType, time: float) -> None: self.translational_constraint_selector, ) - def constrain_rates(self, system: SystemType, time: float) -> None: + def constrain_rates(self, system: SystemType, time: np.floating) -> None: if self.constrained_position_idx.size: self.nb_constrain_translational_rates( system.velocity_collection, @@ -358,7 +367,10 @@ def constrain_rates(self, system: SystemType, time: float) -> None: @staticmethod @njit(cache=True) def nb_constrain_translational_values( - position_collection, fixed_position_collection, indices, constraint_selector + position_collection: NDArray[np.floating], + fixed_position_collection: NDArray[np.floating], + indices: NDArray[np.integer], + constraint_selector: NDArray[np.integer], ) -> None: """ Computes constrain values in numba njit decorator @@ -393,7 +405,9 @@ def nb_constrain_translational_values( @staticmethod @njit(cache=True) def nb_constrain_translational_rates( - velocity_collection, indices, constraint_selector + velocity_collection: NDArray[np.floating], + indices: NDArray[np.integer], + constraint_selector: NDArray[np.integer], ) -> None: """ Compute constrain rates in numba njit decorator @@ -422,7 +436,10 @@ def nb_constrain_translational_rates( @staticmethod @njit(cache=True) def nb_constrain_rotational_rates( - director_collection, omega_collection, indices, constraint_selector + director_collection: NDArray[np.floating], + omega_collection: NDArray[np.floating], + indices: NDArray[np.integer], + constraint_selector: NDArray[np.integer], ) -> None: """ Compute constrain rates in numba njit decorator @@ -489,7 +506,7 @@ class FixedConstraint(GeneralConstraint): GeneralConstraint: Generalized constraint with configurable DOF. """ - def __init__(self, *args, **kwargs): + def __init__(self, *args: Any, **kwargs: Any) -> None: """ Initialization of the constraint. Any parameter passed to 'using' will be available in kwargs. @@ -508,7 +525,7 @@ def __init__(self, *args, **kwargs): **kwargs, ) - def constrain_values(self, system: SystemType, time: float) -> None: + def constrain_values(self, system: SystemType, time: np.floating) -> None: if self.constrained_position_idx.size: self.nb_constrain_translational_values( system.position_collection, @@ -522,7 +539,7 @@ def constrain_values(self, system: SystemType, time: float) -> None: self.constrained_director_idx, ) - def constrain_rates(self, system: SystemType, time: float) -> None: + def constrain_rates(self, system: SystemType, time: np.floating) -> None: if self.constrained_position_idx.size: self.nb_constrain_translational_rates( system.velocity_collection, @@ -537,7 +554,9 @@ def constrain_rates(self, system: SystemType, time: float) -> None: @staticmethod @njit(cache=True) def nb_constraint_rotational_values( - director_collection, fixed_director_collection, indices + director_collection: NDArray[np.floating], + fixed_director_collection: NDArray[np.floating], + indices: NDArray[np.integer], ) -> None: """ Computes constrain values in numba njit decorator @@ -558,7 +577,9 @@ def nb_constraint_rotational_values( @staticmethod @njit(cache=True) def nb_constrain_translational_values( - position_collection, fixed_position_collection, indices + position_collection: NDArray[np.floating], + fixed_position_collection: NDArray[np.floating], + indices: NDArray[np.integer], ) -> None: """ Computes constrain values in numba njit decorator @@ -578,7 +599,9 @@ def nb_constrain_translational_values( @staticmethod @njit(cache=True) - def nb_constrain_translational_rates(velocity_collection, indices) -> None: + def nb_constrain_translational_rates( + velocity_collection: NDArray[np.floating], indices: NDArray[np.integer] + ) -> None: """ Compute constrain rates in numba njit decorator Parameters @@ -598,7 +621,9 @@ def nb_constrain_translational_rates(velocity_collection, indices) -> None: @staticmethod @njit(cache=True) - def nb_constrain_rotational_rates(omega_collection, indices) -> None: + def nb_constrain_rotational_rates( + omega_collection: NDArray[np.floating], indices: NDArray[np.integer] + ) -> None: """ Compute constrain rates in numba njit decorator Parameters @@ -654,15 +679,15 @@ class HelicalBucklingBC(ConstraintBase): def __init__( self, - position_start: np.ndarray, - position_end: np.ndarray, - director_start: np.ndarray, - director_end: np.ndarray, - twisting_time: float, - slack: float, - number_of_rotations: float, - **kwargs, - ): + position_start: NDArray[np.floating], + position_end: NDArray[np.floating], + director_start: NDArray[np.floating], + director_end: NDArray[np.floating], + twisting_time: np.floating, + slack: np.floating, + number_of_rotations: np.floating, + **kwargs: Any, + ) -> None: """ Helical Buckling initializer @@ -718,7 +743,7 @@ def __init__( @ director_end ) # rotation_matrix wants vectors 3,1 - def constrain_values(self, rod: RodType, time: float) -> None: + def constrain_values(self, rod: SystemType, time: np.floating) -> None: if time > self.twisting_time: rod.position_collection[..., 0] = self.final_start_position rod.position_collection[..., -1] = self.final_end_position @@ -726,7 +751,7 @@ def constrain_values(self, rod: RodType, time: float) -> None: rod.director_collection[..., 0] = self.final_start_directors rod.director_collection[..., -1] = self.final_end_directors - def constrain_rates(self, rod: RodType, time: float) -> None: + def constrain_rates(self, rod: SystemType, time: np.floating) -> None: if time > self.twisting_time: rod.velocity_collection[..., 0] = 0.0 rod.omega_collection[..., 0] = 0.0 diff --git a/elastica/callback_functions.py b/elastica/callback_functions.py index ab865dfd0..ac2059a6d 100644 --- a/elastica/callback_functions.py +++ b/elastica/callback_functions.py @@ -4,9 +4,12 @@ import sys import numpy as np import logging +from typing import Any, Optional from collections import defaultdict +from elastica.typing import RodType, SystemType + class CallBackBaseClass: """ @@ -19,13 +22,15 @@ class CallBackBaseClass: """ - def __init__(self): + def __init__(self) -> None: """ CallBackBaseClass does not need any input parameters. """ pass - def make_callback(self, system, time, current_step: int): + def make_callback( + self, system: SystemType, time: np.floating, current_step: int + ) -> None: """ This method is called every time step. Users can define which parameters are called back and recorded. Also users @@ -59,7 +64,7 @@ class MyCallBack(CallBackBaseClass): Collected callback data is saved in this dictionary. """ - def __init__(self, step_skip: int, callback_params): + def __init__(self, step_skip: int, callback_params: dict) -> None: """ Parameters @@ -73,7 +78,9 @@ def __init__(self, step_skip: int, callback_params): self.sample_every = step_skip self.callback_params = callback_params - def make_callback(self, system, time, current_step: int): + def make_callback( + self, system: SystemType, time: np.floating, current_step: int + ) -> None: if current_step % self.sample_every == 0: @@ -116,8 +123,8 @@ def __init__( directory: str, method: str, initial_file_count: int = 0, - file_save_interval: int = 1e8, - ): + file_save_interval: int = 100_000_000, + ) -> None: """ Parameters ---------- @@ -189,7 +196,9 @@ def __init__( self._pickle = pickle self._ext = "pkl" - def make_callback(self, system, time, current_step: int): + def make_callback( + self, system: SystemType, time: np.floating, current_step: int + ) -> None: """ Parameters @@ -224,7 +233,7 @@ def make_callback(self, system, time, current_step: int): ): self._dump() - def _dump(self, **kwargs): + def _dump(self, **kwargs: Any) -> None: """ Dump dictionary buffer (self.buffer) to a file and clear the buffer. @@ -247,7 +256,7 @@ def _dump(self, **kwargs): self.buffer_size = 0 self.buffer.clear() - def get_last_saved_path(self) -> str: + def get_last_saved_path(self) -> Optional[str]: """ Return last saved file path. If no file has been saved, return None @@ -257,14 +266,14 @@ def get_last_saved_path(self) -> str: else: return self.save_path.format(self.file_count - 1, self._ext) - def close(self): + def close(self) -> None: """ Save residual buffer """ if self.buffer_size: self._dump() - def clear(self): + def clear(self) -> None: """ Alias to `close` """ diff --git a/elastica/contact_forces.py b/elastica/contact_forces.py index 8f9b0ab5e..9e9f5a2de 100644 --- a/elastica/contact_forces.py +++ b/elastica/contact_forces.py @@ -1,5 +1,6 @@ __doc__ = """ Numba implementation module containing contact between rods and rigid bodies and other rods rigid bodies or surfaces.""" +from typing import Optional from elastica.typing import RodType, SystemType, AllowedContactType from elastica.rod import RodBase from elastica.rigidbody import Cylinder, Sphere @@ -19,6 +20,7 @@ _calculate_contact_forces_cylinder_plane, ) import numpy as np +from numpy.typing import NDArray class NoContact: @@ -32,7 +34,7 @@ class NoContact: """ - def __init__(self): + def __init__(self) -> None: """ NoContact class does not need any input parameters. """ @@ -69,7 +71,7 @@ def apply_contact( self, system_one: SystemType, system_two: AllowedContactType, - ) -> None: + ) -> Optional[tuple[NDArray[np.floating], NDArray[np.intp]]]: """ Apply contact forces and torques between SystemType object and AllowedContactType object. @@ -101,7 +103,7 @@ class RodRodContact(NoContact): """ - def __init__(self, k: float, nu: float): + def __init__(self, k: np.floating, nu: np.floating) -> None: """ Parameters ---------- @@ -225,11 +227,11 @@ class RodCylinderContact(NoContact): def __init__( self, - k: float, - nu: float, - velocity_damping_coefficient=0.0, - friction_coefficient=0.0, - ): + k: np.floating, + nu: np.floating, + velocity_damping_coefficient: np.floating = 0.0, + friction_coefficient: np.floating = 0.0, + ) -> None: """ Parameters @@ -338,7 +340,7 @@ class RodSelfContact(NoContact): """ - def __init__(self, k: float, nu: float): + def __init__(self, k: np.floating, nu: np.floating) -> None: """ Parameters @@ -435,11 +437,11 @@ class RodSphereContact(NoContact): def __init__( self, - k: float, - nu: float, - velocity_damping_coefficient=0.0, - friction_coefficient=0.0, - ): + k: np.floating, + nu: np.floating, + velocity_damping_coefficient: np.floating = 0.0, + friction_coefficient: np.floating = 0.0, + ) -> None: """ Parameters ---------- @@ -560,9 +562,9 @@ class RodPlaneContact(NoContact): def __init__( self, - k: float, - nu: float, - ): + k: np.floating, + nu: np.floating, + ) -> None: """ Parameters ---------- @@ -652,12 +654,12 @@ class RodPlaneContactWithAnisotropicFriction(NoContact): def __init__( self, - k: float, - nu: float, - slip_velocity_tol: float, - static_mu_array: np.ndarray, - kinetic_mu_array: np.ndarray, - ): + k: np.floating, + nu: np.floating, + slip_velocity_tol: np.floating, + static_mu_array: NDArray[np.floating], + kinetic_mu_array: NDArray[np.floating], + ) -> None: """ Parameters ---------- @@ -776,9 +778,9 @@ class CylinderPlaneContact(NoContact): def __init__( self, - k: float, - nu: float, - ): + k: np.floating, + nu: np.floating, + ) -> None: """ Parameters ---------- @@ -818,7 +820,9 @@ def _check_systems_validity( ) ) - def apply_contact(self, system_one: Cylinder, system_two: SystemType): + def apply_contact( + self, system_one: Cylinder, system_two: SystemType + ) -> tuple[NDArray[np.floating], NDArray[np.intp]]: """ This function computes the plane force response on the cylinder, in the case of contact. Contact model given in Eqn 4.8 Gazzola et. al. RSoS 2018 paper diff --git a/elastica/contact_utils.py b/elastica/contact_utils.py index 71584d2bb..657b6dd00 100644 --- a/elastica/contact_utils.py +++ b/elastica/contact_utils.py @@ -3,37 +3,47 @@ from math import sqrt import numba import numpy as np +from numpy.typing import NDArray from elastica._linalg import ( _batch_norm, ) +from typing import Literal, Sequence, TypeVar @numba.njit(cache=True) -def _dot_product(a, b): - sum = 0.0 +def _dot_product(a: Sequence[np.floating], b: Sequence[np.floating]) -> np.floating: + sum: np.floating = 0.0 for i in range(3): sum += a[i] * b[i] return sum @numba.njit(cache=True) -def _norm(a): +def _norm(a: Sequence[np.floating]) -> float: return sqrt(_dot_product(a, a)) +_SupportsCompareT = TypeVar("_SupportsCompareT") + + @numba.njit(cache=True) -def _clip(x, low, high): +def _clip(x: np.floating, low: np.floating, high: np.floating) -> np.floating: return max(low, min(x, high)) # Can this be made more efficient than 2 comp, 1 or? @numba.njit(cache=True) -def _out_of_bounds(x, low, high): +def _out_of_bounds(x: np.floating, low: np.floating, high: np.floating) -> bool: return (x < low) or (x > high) @numba.njit(cache=True) -def _find_min_dist(x1, e1, x2, e2): +def _find_min_dist( + x1: NDArray[np.floating], + e1: NDArray[np.floating], + x2: NDArray[np.floating], + e2: NDArray[np.floating], +) -> tuple[NDArray[np.floating], NDArray[np.floating], NDArray[np.floating]]: e1e1 = _dot_product(e1, e1) e1e2 = _dot_product(e1, e2) e2e2 = _dot_product(e2, e2) @@ -99,7 +109,9 @@ def _find_min_dist(x1, e1, x2, e2): @numba.njit(cache=True) -def _aabbs_not_intersecting(aabb_one, aabb_two): +def _aabbs_not_intersecting( + aabb_one: NDArray[np.floating], aabb_two: NDArray[np.floating] +) -> Literal[1, 0]: """Returns true if not intersecting else false""" if (aabb_one[0, 1] < aabb_two[0, 0]) | (aabb_one[0, 0] > aabb_two[0, 1]): return 1 @@ -113,14 +125,14 @@ def _aabbs_not_intersecting(aabb_one, aabb_two): @numba.njit(cache=True) def _prune_using_aabbs_rod_cylinder( - rod_one_position_collection, - rod_one_radius_collection, - rod_one_length_collection, - cylinder_position, - cylinder_director, - cylinder_radius, - cylinder_length, -): + rod_one_position_collection: NDArray[np.floating], + rod_one_radius_collection: NDArray[np.floating], + rod_one_length_collection: NDArray[np.floating], + cylinder_position: NDArray[np.floating], + cylinder_director: NDArray[np.floating], + cylinder_radius: NDArray[np.floating], + cylinder_length: NDArray[np.floating], +) -> Literal[1, 0]: max_possible_dimension = np.zeros((3,)) aabb_rod = np.empty((3, 2)) aabb_cylinder = np.empty((3, 2)) @@ -155,13 +167,13 @@ def _prune_using_aabbs_rod_cylinder( @numba.njit(cache=True) def _prune_using_aabbs_rod_rod( - rod_one_position_collection, - rod_one_radius_collection, - rod_one_length_collection, - rod_two_position_collection, - rod_two_radius_collection, - rod_two_length_collection, -): + rod_one_position_collection: NDArray[np.floating], + rod_one_radius_collection: NDArray[np.floating], + rod_one_length_collection: NDArray[np.floating], + rod_two_position_collection: NDArray[np.floating], + rod_two_radius_collection: NDArray[np.floating], + rod_two_length_collection: NDArray[np.floating], +) -> Literal[1, 0]: max_possible_dimension = np.zeros((3,)) aabb_rod_one = np.empty((3, 2)) aabb_rod_two = np.empty((3, 2)) @@ -193,13 +205,13 @@ def _prune_using_aabbs_rod_rod( @numba.njit(cache=True) def _prune_using_aabbs_rod_sphere( - rod_one_position_collection, - rod_one_radius_collection, - rod_one_length_collection, - sphere_position, - sphere_director, - sphere_radius, -): + rod_one_position_collection: NDArray[np.floating], + rod_one_radius_collection: NDArray[np.floating], + rod_one_length_collection: NDArray[np.floating], + sphere_position: NDArray[np.floating], + sphere_director: NDArray[np.floating], + sphere_radius: NDArray[np.floating], +) -> Literal[1, 0]: max_possible_dimension = np.zeros((3,)) aabb_rod = np.empty((3, 2)) aabb_sphere = np.empty((3, 2)) @@ -231,7 +243,9 @@ def _prune_using_aabbs_rod_sphere( @numba.njit(cache=True) -def _find_slipping_elements(velocity_slip, velocity_threshold): +def _find_slipping_elements( + velocity_slip: NDArray[np.floating], velocity_threshold: np.floating +) -> NDArray[np.floating]: """ This function takes the velocity of elements and checks if they are larger than the threshold velocity. If the velocity of elements is larger than threshold velocity, that means those elements are slipping. @@ -272,7 +286,7 @@ def _find_slipping_elements(velocity_slip, velocity_threshold): @numba.njit(cache=True) -def _node_to_element_mass_or_force(input): +def _node_to_element_mass_or_force(input: NDArray[np.floating]) -> NDArray[np.floating]: """ This function converts the mass/forces on rod nodes to elements, where special treatment is necessary at the ends. @@ -310,7 +324,10 @@ def _node_to_element_mass_or_force(input): @numba.njit(cache=True) -def _elements_to_nodes_inplace(vector_in_element_frame, vector_in_node_frame): +def _elements_to_nodes_inplace( + vector_in_element_frame: NDArray[np.floating], + vector_in_node_frame: NDArray[np.floating], +) -> None: """ Updating nodal forces using the forces computed on elements Parameters @@ -333,7 +350,9 @@ def _elements_to_nodes_inplace(vector_in_element_frame, vector_in_node_frame): @numba.njit(cache=True) -def _node_to_element_position(node_position_collection): +def _node_to_element_position( + node_position_collection: NDArray[np.floating], +) -> NDArray[np.floating]: """ This function computes the position of the elements from the nodal values. @@ -379,7 +398,9 @@ def _node_to_element_position(node_position_collection): @numba.njit(cache=True) -def _node_to_element_velocity(mass, node_velocity_collection): +def _node_to_element_velocity( + mass: NDArray[np.floating], node_velocity_collection: NDArray[np.floating] +) -> NDArray[np.floating]: """ This function computes the velocity of the elements from the nodal values. Uses the velocity of center of mass diff --git a/elastica/dissipation.py b/elastica/dissipation.py index f3629fe76..f0a79f560 100644 --- a/elastica/dissipation.py +++ b/elastica/dissipation.py @@ -5,12 +5,14 @@ """ from abc import ABC, abstractmethod +from typing import Any from elastica.typing import RodType, SystemType from numba import njit import numpy as np +from numpy.typing import NDArray class DamperBase(ABC): @@ -29,7 +31,8 @@ class DamperBase(ABC): _system: SystemType - def __init__(self, *args, **kwargs): + # TODO typing can be made better + def __init__(self, *args: Any, **kwargs: Any) -> None: """Initialize damping module""" try: self._system = kwargs["_system"] @@ -40,7 +43,7 @@ def __init__(self, *args, **kwargs): ) @property - def system(self): # -> SystemType: (Return type is not parsed with sphinx book.) + def system(self) -> SystemType: """ get system (rod or rigid body) reference @@ -52,7 +55,7 @@ def system(self): # -> SystemType: (Return type is not parsed with sphinx book. return self._system @abstractmethod - def dampen_rates(self, system: SystemType, time: float): + def dampen_rates(self, system: SystemType, time: np.floating) -> None: # TODO: In the future, we can remove rod and use self.system """ Dampen rates (velocity and/or omega) of a rod object. @@ -113,7 +116,9 @@ class AnalyticalLinearDamper(DamperBase): Damping coefficient acting on rotational velocity. """ - def __init__(self, damping_constant, time_step, **kwargs): + def __init__( + self, damping_constant: np.floating, time_step: np.floating, **kwargs: Any + ) -> None: """ Analytical linear damper initializer @@ -143,7 +148,7 @@ def __init__(self, damping_constant, time_step, **kwargs): * np.diagonal(self._system.inv_mass_second_moment_of_inertia).T ) - def dampen_rates(self, rod: RodType, time: float): + def dampen_rates(self, rod: SystemType, time: np.floating) -> None: rod.velocity_collection[:] = ( rod.velocity_collection * self.translational_damping_coefficient ) @@ -202,7 +207,7 @@ class LaplaceDissipationFilter(DamperBase): Filter term that modifies rod rotational velocity. """ - def __init__(self, filter_order: int, **kwargs): + def __init__(self, filter_order: int, **kwargs: Any) -> None: """ Filter damper initializer @@ -232,7 +237,7 @@ def __init__(self, filter_order: int, **kwargs): self.omega_filter_term = np.zeros_like(self._system.omega_collection) self.filter_function = _filter_function_periodic_condition - def dampen_rates(self, rod: RodType, time: float) -> None: + def dampen_rates(self, rod: RodType, time: np.floating) -> None: self.filter_function( rod.velocity_collection, @@ -245,12 +250,12 @@ def dampen_rates(self, rod: RodType, time: float) -> None: @njit(cache=True) def _filter_function_periodic_condition_ring_rod( - velocity_collection, - velocity_filter_term, - omega_collection, - omega_filter_term, - filter_order, -): + velocity_collection: NDArray[np.floating], + velocity_filter_term: NDArray[np.floating], + omega_collection: NDArray[np.floating], + omega_filter_term: NDArray[np.floating], + filter_order: int, +) -> None: blocksize = velocity_filter_term.shape[1] # Transfer velocity to an array which has periodic boundaries and synchornize boundaries @@ -283,12 +288,12 @@ def _filter_function_periodic_condition_ring_rod( @njit(cache=True) def _filter_function_periodic_condition( - velocity_collection, - velocity_filter_term, - omega_collection, - omega_filter_term, - filter_order, -): + velocity_collection: NDArray[np.floating], + velocity_filter_term: NDArray[np.floating], + omega_collection: NDArray[np.floating], + omega_filter_term: NDArray[np.floating], + filter_order: int, +) -> None: nb_filter_rate( rate_collection=velocity_collection, filter_term=velocity_filter_term, @@ -303,7 +308,9 @@ def _filter_function_periodic_condition( @njit(cache=True) def nb_filter_rate( - rate_collection: np.ndarray, filter_term: np.ndarray, filter_order: int + rate_collection: NDArray[np.floating], + filter_term: NDArray[np.floating], + filter_order: int, ) -> None: """ Filters the rod rates (velocities) in numba njit decorator diff --git a/elastica/external_forces.py b/elastica/external_forces.py index cb9c61e6c..62c4e6adb 100644 --- a/elastica/external_forces.py +++ b/elastica/external_forces.py @@ -3,6 +3,8 @@ import numpy as np +from numpy.typing import NDArray + from elastica._linalg import _batch_matvec from elastica.typing import SystemType, RodType from elastica.utils import _bspline @@ -22,13 +24,13 @@ class NoForces: """ - def __init__(self): + def __init__(self) -> None: """ NoForces class does not need any input parameters. """ pass - def apply_forces(self, system: SystemType, time: np.float64 = 0.0): + def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: """Apply forces to a rod-like object. In NoForces class, this routine simply passes. @@ -43,7 +45,7 @@ def apply_forces(self, system: SystemType, time: np.float64 = 0.0): """ pass - def apply_torques(self, system: SystemType, time: np.float64 = 0.0): + def apply_torques(self, system: SystemType, time: np.floating = 0.0) -> None: """Apply torques to a rod-like object. In NoForces class, this routine simply passes. @@ -70,7 +72,9 @@ class GravityForces(NoForces): """ - def __init__(self, acc_gravity=np.array([0.0, -9.80665, 0.0])): + def __init__( + self, acc_gravity: NDArray[np.floating] = np.array([0.0, -9.80665, 0.0]) + ) -> None: """ Parameters @@ -82,14 +86,18 @@ def __init__(self, acc_gravity=np.array([0.0, -9.80665, 0.0])): super(GravityForces, self).__init__() self.acc_gravity = acc_gravity - def apply_forces(self, system: SystemType, time=0.0): + def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: self.compute_gravity_forces( self.acc_gravity, system.mass, system.external_forces ) @staticmethod @njit(cache=True) - def compute_gravity_forces(acc_gravity, mass, external_forces): + def compute_gravity_forces( + acc_gravity: NDArray[np.floating], + mass: NDArray[np.floating], + external_forces: NDArray[np.floating], + ) -> None: """ This function add gravitational forces on the nodes. We are using njit decorated function to increase the speed. @@ -122,7 +130,12 @@ class EndpointForces(NoForces): """ - def __init__(self, start_force, end_force, ramp_up_time): + def __init__( + self, + start_force: NDArray[np.floating], + end_force: NDArray[np.floating], + ramp_up_time: np.floating, + ) -> None: """ Parameters @@ -143,7 +156,7 @@ def __init__(self, start_force, end_force, ramp_up_time): assert ramp_up_time > 0.0 self.ramp_up_time = ramp_up_time - def apply_forces(self, system: SystemType, time=0.0): + def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: self.compute_end_point_forces( system.external_forces, self.start_force, @@ -155,8 +168,12 @@ def apply_forces(self, system: SystemType, time=0.0): @staticmethod @njit(cache=True) def compute_end_point_forces( - external_forces, start_force, end_force, time, ramp_up_time - ): + external_forces: NDArray[np.floating], + start_force: NDArray[np.floating], + end_force: NDArray[np.floating], + time: np.floating, + ramp_up_time: np.floating, + ) -> None: """ Compute end point forces that are applied on the rod using numba njit decorator. @@ -174,7 +191,7 @@ def compute_end_point_forces( Applied forces are ramped up until ramp up time. """ - factor = min(1.0, time / ramp_up_time) + factor: np.floating = min(1.0, time / ramp_up_time) external_forces[..., 0] += start_force * factor external_forces[..., -1] += end_force * factor @@ -190,7 +207,11 @@ class UniformTorques(NoForces): """ - def __init__(self, torque, direction=np.array([0.0, 0.0, 0.0])): + def __init__( + self, + torque: np.floating, + direction: NDArray[np.floating] = np.array([0.0, 0.0, 0.0]), + ) -> None: """ Parameters @@ -204,7 +225,7 @@ def __init__(self, torque, direction=np.array([0.0, 0.0, 0.0])): super(UniformTorques, self).__init__() self.torque = torque * direction - def apply_torques(self, system: SystemType, time: np.float64 = 0.0): + def apply_torques(self, system: SystemType, time: np.floating = 0.0) -> None: n_elems = system.n_elems torque_on_one_element = ( _batch_product_i_k_to_ik(self.torque, np.ones((n_elems))) / n_elems @@ -224,7 +245,11 @@ class UniformForces(NoForces): 2D (dim, 1) array containing data with 'float' type. Total force applied to a rod-like object. """ - def __init__(self, force, direction=np.array([0.0, 0.0, 0.0])): + def __init__( + self, + force: np.floating, + direction: NDArray[np.floating] = np.array([0.0, 0.0, 0.0]), + ) -> None: """ Parameters @@ -238,7 +263,7 @@ def __init__(self, force, direction=np.array([0.0, 0.0, 0.0])): super(UniformForces, self).__init__() self.force = (force * direction).reshape(3, 1) - def apply_forces(self, rod: RodType, time: np.float64 = 0.0): + def apply_forces(self, rod: SystemType, time: np.floating = 0.0) -> None: force_on_one_element = self.force / rod.n_elems rod.external_forces += force_on_one_element @@ -275,16 +300,16 @@ class MuscleTorques(NoForces): def __init__( self, - base_length, - b_coeff, - period, - wave_number, - phase_shift, - direction, - rest_lengths, - ramp_up_time, - with_spline=False, - ): + base_length: np.floating, + b_coeff: NDArray[np.floating], + period: np.floating, + wave_number: np.floating, + phase_shift: np.floating, + direction: NDArray[np.floating], + rest_lengths: NDArray[np.floating], + ramp_up_time: np.floating, + with_spline: bool = False, + ) -> None: """ Parameters @@ -335,7 +360,7 @@ def __init__( else: self.my_spline = np.full_like(self.s, fill_value=1.0) - def apply_torques(self, rod: RodType, time: np.float64 = 0.0): + def apply_torques(self, rod: SystemType, time: np.floating = 0.0) -> None: self.compute_muscle_torques( time, self.my_spline, @@ -352,19 +377,19 @@ def apply_torques(self, rod: RodType, time: np.float64 = 0.0): @staticmethod @njit(cache=True) def compute_muscle_torques( - time, - my_spline, - s, - angular_frequency, - wave_number, - phase_shift, - ramp_up_time, - direction, - director_collection, - external_torques, - ): + time: np.floating, + my_spline: NDArray[np.floating], + s: np.floating, + angular_frequency: np.floating, + wave_number: np.floating, + phase_shift: np.floating, + ramp_up_time: np.floating, + direction: NDArray[np.floating], + director_collection: NDArray[np.floating], + external_torques: NDArray[np.floating], + ) -> None: # Ramp up the muscle torque - factor = min(1.0, time / ramp_up_time) + factor: np.floating = min(1.0, time / ramp_up_time) # From the node 1 to node nelem-1 # Magnitude of the torque. Am = beta(s) * sin(2pi*t/T + 2pi*s/lambda + phi) # There is an inconsistency with paper and Elastica cpp implementation. In paper sign in @@ -388,7 +413,10 @@ def compute_muscle_torques( @njit(cache=True) -def inplace_addition(external_force_or_torque, force_or_torque): +def inplace_addition( + external_force_or_torque: NDArray[np.floating], + force_or_torque: NDArray[np.floating], +) -> None: """ This function does inplace addition. First argument `external_force_or_torque` is the system.external_forces @@ -411,7 +439,10 @@ def inplace_addition(external_force_or_torque, force_or_torque): @njit(cache=True) -def inplace_substraction(external_force_or_torque, force_or_torque): +def inplace_substraction( + external_force_or_torque: NDArray[np.floating], + force_or_torque: NDArray[np.floating], +) -> None: """ This function does inplace substraction. First argument `external_force_or_torque` is the system.external_forces @@ -460,12 +491,12 @@ class EndpointForcesSinusoidal(NoForces): def __init__( self, - start_force_mag, - end_force_mag, - ramp_up_time=0.0, - tangent_direction=np.array([0, 0, 1]), - normal_direction=np.array([0, 1, 0]), - ): + start_force_mag: np.floating, + end_force_mag: np.floating, + ramp_up_time: np.floating = 0.0, + tangent_direction: NDArray[np.floating] = np.array([0, 0, 1]), + normal_direction: NDArray[np.floating] = np.array([0, 1, 0]), + ) -> None: """ Parameters @@ -495,7 +526,7 @@ def __init__( assert ramp_up_time >= 0.0 self.ramp_up_time = ramp_up_time - def apply_forces(self, system: SystemType, time=0.0): + def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: if time < self.ramp_up_time: # When time smaller than ramp up time apply the force in normal direction diff --git a/elastica/interaction.py b/elastica/interaction.py index 95d4602e1..d695e6153 100644 --- a/elastica/interaction.py +++ b/elastica/interaction.py @@ -1,6 +1,7 @@ __doc__ = """ Numba implementation module containing interactions between a rod and its environment.""" +from typing import Any, NoReturn import numpy as np from elastica.external_forces import NoForces from numba import njit @@ -14,8 +15,12 @@ _calculate_contact_forces_cylinder_plane, ) +from numpy.typing import NDArray -def find_slipping_elements(velocity_slip, velocity_threshold): +from elastica.typing import SystemType + + +def find_slipping_elements(velocity_slip: Any, velocity_threshold: Any) -> NoReturn: raise NotImplementedError( "This function is removed in v0.3.2. Please use\n" "elastica.contact_utils._find_slipping_elements()\n" @@ -23,7 +28,7 @@ def find_slipping_elements(velocity_slip, velocity_threshold): ) -def node_to_element_mass_or_force(input): +def node_to_element_mass_or_force(input: Any) -> NoReturn: raise NotImplementedError( "This function is removed in v0.3.2. Please use\n" "elastica.contact_utils._node_to_element_mass_or_force()\n" @@ -31,7 +36,7 @@ def node_to_element_mass_or_force(input): ) -def nodes_to_elements(input): +def nodes_to_elements(input: Any) -> NoReturn: # Remove the function beyond v0.4.0 raise NotImplementedError( "This function is removed in v0.3.1. Please use\n" @@ -41,7 +46,9 @@ def nodes_to_elements(input): @njit(cache=True) -def elements_to_nodes_inplace(vector_in_element_frame, vector_in_node_frame): +def elements_to_nodes_inplace( + vector_in_element_frame: Any, vector_in_node_frame: Any +) -> NoReturn: raise NotImplementedError( "This function is removed in v0.3.2. Please use\n" "elastica.contact_utils._elements_to_nodes_inplace()\n" @@ -74,7 +81,13 @@ class InteractionPlane: """ - def __init__(self, k, nu, plane_origin, plane_normal): + def __init__( + self, + k: np.floating, + nu: np.floating, + plane_origin: NDArray[np.floating], + plane_normal: NDArray[np.floating], + ) -> None: """ Parameters @@ -96,7 +109,9 @@ def __init__(self, k, nu, plane_origin, plane_normal): self.plane_normal = plane_normal.reshape(3) self.surface_tol = 1e-4 - def apply_normal_force(self, system): + def apply_normal_force( + self, system: SystemType + ) -> tuple[NDArray[np.floating], NDArray[np.intp]]: """ In the case of contact with the plane, this function computes the plane reaction force on the element. @@ -130,18 +145,18 @@ def apply_normal_force(self, system): def apply_normal_force_numba( - plane_origin, - plane_normal, - surface_tol, - k, - nu, - radius, - mass, - position_collection, - velocity_collection, - internal_forces, - external_forces, -): + plane_origin: Any, + plane_normal: Any, + surface_tol: Any, + k: Any, + nu: Any, + radius: Any, + mass: Any, + position_collection: Any, + velocity_collection: Any, + internal_forces: Any, + external_forces: Any, +) -> NoReturn: raise NotImplementedError( "This function is removed in v0.3.2. For rod plane contact please use: \n" "elastica._contact_functions._calculate_contact_forces_rod_plane() \n" @@ -186,14 +201,14 @@ class AnisotropicFrictionalPlane(NoForces, InteractionPlane): def __init__( self, - k, - nu, - plane_origin, - plane_normal, - slip_velocity_tol, - static_mu_array, - kinetic_mu_array, - ): + k: np.floating, + nu: np.floating, + plane_origin: NDArray[np.floating], + plane_normal: NDArray[np.floating], + slip_velocity_tol: np.floating, + static_mu_array: NDArray[np.floating], + kinetic_mu_array: NDArray[np.floating], + ) -> None: """ Parameters @@ -232,7 +247,7 @@ def __init__( # kinetic and static friction should separate functions # for now putting them together to figure out common variables - def apply_forces(self, system, time=0.0): + def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: """ Call numba implementation to apply friction forces Parameters @@ -269,30 +284,30 @@ def apply_forces(self, system, time=0.0): def anisotropic_friction( - plane_origin, - plane_normal, - surface_tol, - slip_velocity_tol, - k, - nu, - kinetic_mu_forward, - kinetic_mu_backward, - kinetic_mu_sideways, - static_mu_forward, - static_mu_backward, - static_mu_sideways, - radius, - mass, - tangents, - position_collection, - director_collection, - velocity_collection, - omega_collection, - internal_forces, - external_forces, - internal_torques, - external_torques, -): + plane_origin: Any, + plane_normal: Any, + surface_tol: Any, + slip_velocity_tol: Any, + k: Any, + nu: Any, + kinetic_mu_forward: Any, + kinetic_mu_backward: Any, + kinetic_mu_sideways: Any, + static_mu_forward: Any, + static_mu_backward: Any, + static_mu_sideways: Any, + radius: Any, + mass: Any, + tangents: Any, + position_collection: Any, + director_collection: Any, + velocity_collection: Any, + omega_collection: Any, + internal_forces: Any, + external_forces: Any, + internal_torques: Any, + external_torques: Any, +) -> NoReturn: raise NotImplementedError( "This function is removed in v0.3.2. For anisotropic_friction please use: \n" "elastica._contact_functions._calculate_contact_forces_rod_plane_with_anisotropic_friction() \n" @@ -302,7 +317,7 @@ def anisotropic_friction( # Slender body module @njit(cache=True) -def sum_over_elements(input): +def sum_over_elements(input: NDArray[np.floating]) -> np.floating: """ This function sums all elements of the input array. Using a Numba njit decorator shows better performance @@ -334,14 +349,14 @@ def sum_over_elements(input): This version: 513 ns ± 24.6 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each) """ - output = 0.0 + output: np.floating = 0.0 for i in range(input.shape[0]): output += input[i] return output -def node_to_element_position(node_position_collection): +def node_to_element_position(node_position_collection: Any) -> NoReturn: raise NotImplementedError( "This function is removed in v0.3.2. For node-to-element_position() interpolation please use: \n" "elastica.contact_utils._node_to_element_position() for rod position \n" @@ -349,7 +364,7 @@ def node_to_element_position(node_position_collection): ) -def node_to_element_velocity(mass, node_velocity_collection): +def node_to_element_velocity(mass: Any, node_velocity_collection: Any) -> NoReturn: raise NotImplementedError( "This function is removed in v0.3.2. For node-to-element_velocity() interpolation please use: \n" "elastica.contact_utils._node_to_element_velocity() for rod velocity. \n" @@ -357,7 +372,7 @@ def node_to_element_velocity(mass, node_velocity_collection): ) -def node_to_element_pos_or_vel(vector_in_node_frame): +def node_to_element_pos_or_vel(vector_in_node_frame: Any) -> NoReturn: # Remove the function beyond v0.4.0 raise NotImplementedError( "This function is removed in v0.3.0. For node-to-element interpolation please use: \n" @@ -369,8 +384,13 @@ def node_to_element_pos_or_vel(vector_in_node_frame): @njit(cache=True) def slender_body_forces( - tangents, velocity_collection, dynamic_viscosity, lengths, radius, mass -): + tangents: NDArray[np.floating], + velocity_collection: NDArray[np.floating], + dynamic_viscosity: np.floating, + lengths: NDArray[np.floating], + radius: NDArray[np.floating], + mass: NDArray[np.floating], +) -> NDArray[np.floating]: r""" This function computes hydrodynamic forces on a body using slender body theory. The below implementation is from Eq. 4.13 in Gazzola et al. RSoS. (2018). @@ -481,7 +501,7 @@ class SlenderBodyTheory(NoForces): """ - def __init__(self, dynamic_viscosity): + def __init__(self, dynamic_viscosity: np.floating) -> None: """ Parameters @@ -492,7 +512,7 @@ def __init__(self, dynamic_viscosity): super(SlenderBodyTheory, self).__init__() self.dynamic_viscosity = dynamic_viscosity - def apply_forces(self, system, time=0.0): + def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: """ This function applies hydrodynamic forces on body using the slender body theory given in @@ -518,14 +538,22 @@ def apply_forces(self, system, time=0.0): # base class for interaction # only applies normal force no friction class InteractionPlaneRigidBody: - def __init__(self, k, nu, plane_origin, plane_normal): + def __init__( + self, + k: np.floating, + nu: np.floating, + plane_origin: NDArray[np.floating], + plane_normal: NDArray[np.floating], + ) -> None: self.k = k self.nu = nu self.plane_origin = plane_origin.reshape(3, 1) self.plane_normal = plane_normal.reshape(3) self.surface_tol = 1e-4 - def apply_normal_force(self, system): + def apply_normal_force( + self, system: SystemType + ) -> tuple[NDArray[np.floating], NDArray[np.intp]]: """ This function computes the plane force response on the rigid body, in the case of contact. Contact model given in Eqn 4.8 Gazzola et. al. RSoS 2018 paper @@ -553,16 +581,16 @@ def apply_normal_force(self, system): @njit(cache=True) def apply_normal_force_numba_rigid_body( - plane_origin, - plane_normal, - surface_tol, - k, - nu, - length, - position_collection, - velocity_collection, - external_forces, -): + plane_origin: Any, + plane_normal: Any, + surface_tol: Any, + k: Any, + nu: Any, + length: Any, + position_collection: Any, + velocity_collection: Any, + external_forces: Any, +) -> NoReturn: raise NotImplementedError( "This function is removed in v0.3.2. For cylinder plane contact please use: \n" diff --git a/elastica/joint.py b/elastica/joint.py index e37d6f058..283783a62 100644 --- a/elastica/joint.py +++ b/elastica/joint.py @@ -1,8 +1,12 @@ __doc__ = """ Module containing joint classes to connect multiple rods together. """ + +from typing import Any, NoReturn, Optional + from elastica._rotations import _inv_rotate from elastica.typing import SystemType, RodType import numpy as np import logging +from numpy.typing import NDArray class FreeJoint: @@ -27,7 +31,7 @@ class FreeJoint: # pass the k and nu for the forces # also the necessary rods for the joint # indices should be 0 or -1, we will provide wrappers for users later - def __init__(self, k, nu): + def __init__(self, k: np.floating, nu: np.floating) -> None: """ Parameters @@ -42,8 +46,12 @@ def __init__(self, k, nu): self.nu = nu def apply_forces( - self, system_one: SystemType, index_one, system_two: SystemType, index_two - ): + self, + system_one: SystemType, + index_one: int, + system_two: SystemType, + index_two: int, + ) -> None: """ Apply joint force to the connected rod objects. @@ -81,8 +89,12 @@ def apply_forces( return def apply_torques( - self, system_one: SystemType, index_one, system_two: SystemType, index_two - ): + self, + system_one: SystemType, + index_one: int, + system_two: SystemType, + index_two: int, + ) -> None: """ Apply restoring joint torques to the connected rod objects. @@ -127,7 +139,13 @@ class HingeJoint(FreeJoint): """ # TODO: IN WRAPPER COMPUTE THE NORMAL DIRECTION OR ASK USER TO GIVE INPUT, IF NOT THROW ERROR - def __init__(self, k, nu, kt, normal_direction): + def __init__( + self, + k: np.floating, + nu: np.floating, + kt: np.floating, + normal_direction: NDArray[np.floating], + ) -> None: """ Parameters @@ -154,19 +172,19 @@ def __init__(self, k, nu, kt, normal_direction): def apply_forces( self, system_one: SystemType, - index_one, + index_one: int, system_two: SystemType, - index_two, - ): + index_two: int, + ) -> None: return super().apply_forces(system_one, index_one, system_two, index_two) def apply_torques( self, system_one: SystemType, - index_one, + index_one: int, system_two: SystemType, - index_two, - ): + index_two: int, + ) -> None: # current tangent direction of the `index_two` element of system two system_two_tangent = system_two.director_collection[2, :, index_two] @@ -215,7 +233,14 @@ class FixedJoint(FreeJoint): is enforced. """ - def __init__(self, k, nu, kt, nut=0.0, rest_rotation_matrix=None): + def __init__( + self, + k: np.floating, + nu: np.floating, + kt: np.floating, + nut: np.floating = 0.0, + rest_rotation_matrix: Optional[NDArray[np.floating]] = None, + ) -> None: """ Parameters @@ -254,19 +279,19 @@ def __init__(self, k, nu, kt, nut=0.0, rest_rotation_matrix=None): def apply_forces( self, system_one: SystemType, - index_one, + index_one: int, system_two: SystemType, - index_two, - ): + index_two: int, + ) -> None: return super().apply_forces(system_one, index_one, system_two, index_two) def apply_torques( self, system_one: SystemType, - index_one, + index_one: int, system_two: SystemType, - index_two, - ): + index_two: int, + ) -> None: # collect directors of systems one and two # note that systems can be either rods or rigid bodies system_one_director = system_one.director_collection[..., index_one] @@ -311,10 +336,10 @@ def apply_torques( def get_relative_rotation_two_systems( system_one: SystemType, - index_one, + index_one: int, system_two: SystemType, - index_two, -): + index_two: int, +) -> NDArray[np.floating]: """ Compute the relative rotation matrix C_12 between system one and system two at the specified elements. @@ -362,7 +387,7 @@ def get_relative_rotation_two_systems( # everything below this comment should be removed beyond v0.4.0 -def _dot_product(a, b): +def _dot_product(a: Any, b: Any) -> NoReturn: raise NotImplementedError( "This function is removed in v0.3.2. Please use\n" "elastica.contact_utils._dot_product()\n" @@ -370,7 +395,7 @@ def _dot_product(a, b): ) -def _norm(a): +def _norm(a: Any) -> NoReturn: raise NotImplementedError( "This function is removed in v0.3.2. Please use\n" "elastica.contact_utils._norm()\n" @@ -378,7 +403,7 @@ def _norm(a): ) -def _clip(x, low, high): +def _clip(x: Any, low: Any, high: Any) -> NoReturn: raise NotImplementedError( "This function is removed in v0.3.2. Please use\n" "elastica.contact_utils._clip()\n" @@ -386,7 +411,7 @@ def _clip(x, low, high): ) -def _out_of_bounds(x, low, high): +def _out_of_bounds(x: Any, low: Any, high: Any) -> NoReturn: raise NotImplementedError( "This function is removed in v0.3.2. Please use\n" "elastica.contact_utils._out_of_bounds()\n" @@ -394,7 +419,7 @@ def _out_of_bounds(x, low, high): ) -def _find_min_dist(x1, e1, x2, e2): +def _find_min_dist(x1: Any, e1: Any, x2: Any, e2: Any) -> NoReturn: raise NotImplementedError( "This function is removed in v0.3.2. Please use\n" "elastica.contact_utils._find_min_dist()\n" @@ -403,25 +428,25 @@ def _find_min_dist(x1, e1, x2, e2): def _calculate_contact_forces_rod_rigid_body( - x_collection_rod, - edge_collection_rod, - x_cylinder_center, - x_cylinder_tip, - edge_cylinder, - radii_sum, - length_sum, - internal_forces_rod, - external_forces_rod, - external_forces_cylinder, - external_torques_cylinder, - cylinder_director_collection, - velocity_rod, - velocity_cylinder, - contact_k, - contact_nu, - velocity_damping_coefficient, - friction_coefficient, -): + x_collection_rod: Any, + edge_collection_rod: Any, + x_cylinder_center: Any, + x_cylinder_tip: Any, + edge_cylinder: Any, + radii_sum: Any, + length_sum: Any, + internal_forces_rod: Any, + external_forces_rod: Any, + external_forces_cylinder: Any, + external_torques_cylinder: Any, + cylinder_director_collection: Any, + velocity_rod: Any, + velocity_cylinder: Any, + contact_k: Any, + contact_nu: Any, + velocity_damping_coefficient: Any, + friction_coefficient: Any, +) -> NoReturn: raise NotImplementedError( "This function is removed in v0.3.2. Please use\n" "elastica._contact_functions._calculate_contact_forces_rod_cylinder()\n" @@ -430,23 +455,23 @@ def _calculate_contact_forces_rod_rigid_body( def _calculate_contact_forces_rod_rod( - x_collection_rod_one, - radius_rod_one, - length_rod_one, - tangent_rod_one, - velocity_rod_one, - internal_forces_rod_one, - external_forces_rod_one, - x_collection_rod_two, - radius_rod_two, - length_rod_two, - tangent_rod_two, - velocity_rod_two, - internal_forces_rod_two, - external_forces_rod_two, - contact_k, - contact_nu, -): + x_collection_rod_one: Any, + radius_rod_one: Any, + length_rod_one: Any, + tangent_rod_one: Any, + velocity_rod_one: Any, + internal_forces_rod_one: Any, + external_forces_rod_one: Any, + x_collection_rod_two: Any, + radius_rod_two: Any, + length_rod_two: Any, + tangent_rod_two: Any, + velocity_rod_two: Any, + internal_forces_rod_two: Any, + external_forces_rod_two: Any, + contact_k: Any, + contact_nu: Any, +) -> NoReturn: raise NotImplementedError( "This function is removed in v0.3.2. Please use\n" "elastica._contact_functions._calculate_contact_forces_rod_rod()\n" @@ -455,15 +480,15 @@ def _calculate_contact_forces_rod_rod( def _calculate_contact_forces_self_rod( - x_collection_rod, - radius_rod, - length_rod, - tangent_rod, - velocity_rod, - external_forces_rod, - contact_k, - contact_nu, -): + x_collection_rod: Any, + radius_rod: Any, + length_rod: Any, + tangent_rod: Any, + velocity_rod: Any, + external_forces_rod: Any, + contact_k: Any, + contact_nu: Any, +) -> NoReturn: raise NotImplementedError( "This function is removed in v0.3.2. Please use\n" "elastica._contact_functions._calculate_contact_forces_self_rod()\n" @@ -471,7 +496,7 @@ def _calculate_contact_forces_self_rod( ) -def _aabbs_not_intersecting(aabb_one, aabb_two): +def _aabbs_not_intersecting(aabb_one: Any, aabb_two: Any) -> NoReturn: raise NotImplementedError( "This function is removed in v0.3.2. Please use\n" "elastica.contact_utils._aabbs_not_intersecting()\n" @@ -480,14 +505,14 @@ def _aabbs_not_intersecting(aabb_one, aabb_two): def _prune_using_aabbs_rod_rigid_body( - rod_one_position_collection, - rod_one_radius_collection, - rod_one_length_collection, - cylinder_position, - cylinder_director, - cylinder_radius, - cylinder_length, -): + rod_one_position_collection: Any, + rod_one_radius_collection: Any, + rod_one_length_collection: Any, + cylinder_position: Any, + cylinder_director: Any, + cylinder_radius: Any, + cylinder_length: Any, +) -> NoReturn: raise NotImplementedError( "This function is removed in v0.3.2. Please use\n" "elastica.contact_utils._prune_using_aabbs_rod_cylinder()\n" @@ -496,13 +521,13 @@ def _prune_using_aabbs_rod_rigid_body( def _prune_using_aabbs_rod_rod( - rod_one_position_collection, - rod_one_radius_collection, - rod_one_length_collection, - rod_two_position_collection, - rod_two_radius_collection, - rod_two_length_collection, -): + rod_one_position_collection: Any, + rod_one_radius_collection: Any, + rod_one_length_collection: Any, + rod_two_position_collection: Any, + rod_two_radius_collection: Any, + rod_two_length_collection: Any, +) -> NoReturn: raise NotImplementedError( "This function is removed in v0.3.2. Please use\n" "elastica.contact_utils._prune_using_aabbs_rod_rod()\n" @@ -555,7 +580,13 @@ class ExternalContact(FreeJoint): # potentially dangerous as it does not deal with "end" conditions # correctly. - def __init__(self, k, nu, velocity_damping_coefficient=0, friction_coefficient=0): + def __init__( + self, + k: np.floating, + nu: np.floating, + velocity_damping_coefficient: np.floating = 0, + friction_coefficient: np.floating = 0, + ) -> None: """ Parameters @@ -586,11 +617,11 @@ def __init__(self, k, nu, velocity_damping_coefficient=0, friction_coefficient=0 def apply_forces( self, - rod_one: RodType, - index_one, + rod_one: SystemType, + index_one: int, rod_two: SystemType, - index_two, - ): + index_two: int, + ) -> None: # del index_one, index_two from elastica.contact_utils import ( _prune_using_aabbs_rod_cylinder, @@ -693,7 +724,7 @@ class SelfContact(FreeJoint): """ - def __init__(self, k, nu): + def __init__(self, k: np.floating, nu: np.floating) -> None: super().__init__(k, nu) log = logging.getLogger(self.__class__.__name__) log.warning( @@ -705,7 +736,9 @@ def __init__(self, k, nu): "The option to use the SelfContact joint for the rod self contact will be removed in the future (v0.3.3).\n" ) - def apply_forces(self, rod_one: RodType, index_one, rod_two: SystemType, index_two): + def apply_forces( + self, rod_one: SystemType, index_one: int, rod_two: SystemType, index_two: int + ) -> None: # del index_one, index_two from elastica._contact_functions import ( _calculate_contact_forces_self_rod, diff --git a/elastica/restart.py b/elastica/restart.py index 1b5fab267..643e0ed03 100644 --- a/elastica/restart.py +++ b/elastica/restart.py @@ -5,8 +5,10 @@ from itertools import groupby from .memory_block import MemoryBlockCosseratRod, MemoryBlockRigidBody +from typing import Iterable, Iterator, Any -def all_equal(iterable): + +def all_equal(iterable: Iterable[Any]) -> bool: """ Checks if all elements of list are equal. Parameters @@ -20,11 +22,17 @@ def all_equal(iterable): ---------- https://stackoverflow.com/questions/3844801/check-if-all-elements-in-a-list-are-identical """ - g = groupby(iterable) + g: Iterator[Any] = groupby(iterable) return next(g, True) and not next(g, False) -def save_state(simulator, directory: str = "", time=0.0, verbose: bool = False): +# TODO: simulator should have better typing +def save_state( + simulator: Iterable, + directory: str = "", + time: np.floating = 0.0, + verbose: bool = False, +) -> None: """ Save state parameters of each rod. TODO : environment list variable is not uniform at the current stage of development. @@ -53,7 +61,10 @@ def save_state(simulator, directory: str = "", time=0.0, verbose: bool = False): print("Save complete: {}".format(directory)) -def load_state(simulator, directory: str = "", verbose: bool = False): +# TODO: simulator should have better typing +def load_state( + simulator: Iterable, directory: str = "", verbose: bool = False +) -> float: """ Load the rod-state. Compatibale with 'save_state' method. If the save-file does not exist, it returns error. @@ -72,7 +83,7 @@ def load_state(simulator, directory: str = "", verbose: bool = False): time : float Simulation time of systems when they are saved. """ - time_list = [] # Simulation time of rods when they are saved. + time_list: list[float] = [] # Simulation time of rods when they are saved. for idx, rod in enumerate(simulator): if isinstance(rod, MemoryBlockCosseratRod) or isinstance( rod, MemoryBlockRigidBody diff --git a/elastica/transformations.py b/elastica/transformations.py index b002e9569..c8e977211 100644 --- a/elastica/transformations.py +++ b/elastica/transformations.py @@ -10,11 +10,14 @@ from .utils import MaxDimension, isqrt +from numpy.typing import NDArray # TODO Complete, but nicer interface, evolve it eventually -def format_vector_shape(vector_collection): +def format_vector_shape( + vector_collection: NDArray[np.floating], +) -> NDArray[np.floating]: """ Function for formatting vector shapes into correct format @@ -59,7 +62,9 @@ def format_vector_shape(vector_collection): return vector_collection -def format_matrix_shape(matrix_collection): +def format_matrix_shape( + matrix_collection: NDArray[np.floating], +) -> NDArray[np.floating]: """ Formats input matrix into correct format @@ -77,7 +82,7 @@ def format_matrix_shape(matrix_collection): # check first two dimensions are same and matrix is square # other possibility is one dimension is dim**2 and other is blocksize, # we need to convert the matrix in that case. - def assert_proper_square(num1): + def assert_proper_square(num1: int) -> int: sqrt_num = isqrt(num1) assert sqrt_num**2 == num1, "Matrix dimension passed is not a perfect square" return sqrt_num @@ -136,12 +141,14 @@ def assert_proper_square(num1): return matrix_collection -def skew_symmetrize(vector): +def skew_symmetrize(vector: NDArray[np.floating]) -> NDArray[np.floating]: vector = format_vector_shape(vector) return _skew_symmetrize(vector) -def inv_skew_symmetrize(matrix_collection): +def inv_skew_symmetrize( + matrix_collection: NDArray[np.floating], +) -> NDArray[np.floating]: """ Safe wrapper around inv_skew_symmetrize that does checking and formatting on type of matrix_collection using format_matrix_shape @@ -167,7 +174,9 @@ def inv_skew_symmetrize(matrix_collection): raise ValueError("matrix_collection passed is not skew-symmetric") -def rotate(matrix, scale, axis): +def rotate( + matrix: NDArray[np.floating], scale: np.floating, axis: NDArray[np.floating] +) -> NDArray[np.floating]: """ This function takes single or multiple frames as matrix. Then rotates these frames around a single axis for all frames, or can rotate each frame around its own diff --git a/elastica/typing.py b/elastica/typing.py index 6599f1af9..fd1337fe5 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -70,7 +70,7 @@ # [SymplecticStepperProtocol, np.floating], np.floating # ] OperatorType: TypeAlias = Callable[ - Any, Any + ..., Any ] # TODO: Maybe can be more specific. Up for discussion. SteppersOperatorsType: TypeAlias = tuple[tuple[OperatorType, ...], ...] # tuple[Union[PrefactorOperatorType, StepOperatorType, NoOpType, np.floating], ...], ... diff --git a/elastica/utils.py b/elastica/utils.py index bbf9baa2d..81517d54d 100644 --- a/elastica/utils.py +++ b/elastica/utils.py @@ -1,12 +1,15 @@ """ Handy utilities """ +from typing import Generator, Iterable, Any, Literal, TypeVar import functools import numpy as np from numpy import finfo, float64 from itertools import islice from scipy.interpolate import BSpline +from numpy.typing import NDArray + # Slower than the python3.8 isqrt implementation for small ints # python isqrt : ~130 ns @@ -47,6 +50,8 @@ def isqrt(num: int) -> int: elif num == 0: return 0 + raise ValueError("num must be a positive number") + class MaxDimension: """ @@ -54,7 +59,7 @@ class MaxDimension: """ @staticmethod - def value(): + def value() -> Literal[3]: """ Returns spatial dimension @@ -67,7 +72,7 @@ def value(): class Tolerance: @staticmethod - def atol(): + def atol() -> np.floating: """ Static absolute tolerance method @@ -78,7 +83,7 @@ def atol(): return finfo(float64).eps * 1e4 @staticmethod - def rtol(): + def rtol() -> np.floating: """ Static relative tolerance method @@ -89,7 +94,7 @@ def rtol(): return finfo(float64).eps * 1e11 -def perm_parity(lst): +def perm_parity(lst: list[int]) -> int: """ Given a permutation of the digits 0..N in order as a list, returns its parity (or sign): +1 for even parity; -1 for odd. @@ -115,7 +120,10 @@ def perm_parity(lst): return parity -def grouper(iterable, n): +_T = TypeVar("_T") + + +def grouper(iterable: Iterable[_T], n: int) -> Generator[tuple[_T, ...], None, None]: """Collect data into fixed-length chunks or blocks" Parameters @@ -144,7 +152,7 @@ def grouper(iterable, n): yield group -def extend_instance(obj, cls): +def extend_instance(obj: Any, cls: Any) -> None: """ Apply mixins to a class instance after creation @@ -170,7 +178,9 @@ def extend_instance(obj, cls): obj.__class__ = type(base_cls_name, (cls, base_cls), {}) -def _bspline(t_coeff, l_centerline=1.0): +def _bspline( + t_coeff: NDArray, l_centerline: np.floating = 1.0 +) -> tuple[BSpline, NDArray, NDArray]: """Generates a bspline object that plots the spline interpolant for any vector x. Optionally takes in a centerline length, set to 1.0 by default and keep_pts for keeping record of control points @@ -198,7 +208,9 @@ def _bspline(t_coeff, l_centerline=1.0): return __bspline_impl__(control_pts, t_coeff, degree) -def __bspline_impl__(x_pts, t_c, degree): +def __bspline_impl__( + x_pts: NDArray, t_c: NDArray, degree: int +) -> tuple[BSpline, NDArray, NDArray]: """""" # Update the knots From b06b8f7945d44470642234be28a483906eebdf78 Mon Sep 17 00:00:00 2001 From: Ankith Date: Sun, 19 May 2024 22:20:23 +0530 Subject: [PATCH 053/134] Improve type hinting in the rod directory --- elastica/rod/cosserat_rod.py | 322 ++++++++++++++++--------------- elastica/rod/data_structures.py | 102 +++++++--- elastica/rod/factory_function.py | 76 ++++++-- elastica/rod/knot_theory.py | 64 +++--- elastica/rod/rod_base.py | 15 +- 5 files changed, 354 insertions(+), 225 deletions(-) diff --git a/elastica/rod/cosserat_rod.py b/elastica/rod/cosserat_rod.py index 051b37513..7d8a4f4af 100644 --- a/elastica/rod/cosserat_rod.py +++ b/elastica/rod/cosserat_rod.py @@ -2,6 +2,7 @@ import numpy as np +from numpy.typing import NDArray import functools import numba from elastica.rod import RodBase @@ -20,18 +21,21 @@ _difference, _average, ) -from typing import Optional +from typing import Any, Optional +from typing_extensions import Self + +from elastica.typing import RodType position_difference_kernel = _difference position_average = _average @functools.lru_cache(maxsize=1) -def _get_z_vector(): +def _get_z_vector() -> NDArray[np.floating]: return np.array([0.0, 0.0, 1.0]).reshape(3, -1) -def _compute_sigma_kappa_for_blockstructure(memory_block): +def _compute_sigma_kappa_for_blockstructure(memory_block: RodType) -> None: """ This function is a wrapper to call functions which computes shear stretch, strain and bending twist and strain. @@ -147,39 +151,39 @@ class CosseratRod(RodBase, KnotTheory): def __init__( self, - n_elements, - position, - velocity, - omega, - acceleration, - angular_acceleration, - directors, - radius, - mass_second_moment_of_inertia, - inv_mass_second_moment_of_inertia, - shear_matrix, - bend_matrix, - density, - volume, - mass, - internal_forces, - internal_torques, - external_forces, - external_torques, - lengths, - rest_lengths, - tangents, - dilatation, - dilatation_rate, - voronoi_dilatation, - rest_voronoi_lengths, - sigma, - kappa, - rest_sigma, - rest_kappa, - internal_stress, - internal_couple, - ring_rod_flag, + n_elements: int, + position: NDArray[np.floating], + velocity: NDArray[np.floating], + omega: NDArray[np.floating], + acceleration: NDArray[np.floating], + angular_acceleration: NDArray[np.floating], + directors: NDArray[np.floating], + radius: NDArray[np.floating], + mass_second_moment_of_inertia: NDArray[np.floating], + inv_mass_second_moment_of_inertia: NDArray[np.floating], + shear_matrix: NDArray[np.floating], + bend_matrix: NDArray[np.floating], + density: NDArray[np.floating], + volume: NDArray[np.floating], + mass: NDArray[np.floating], + internal_forces: NDArray[np.floating], + internal_torques: NDArray[np.floating], + external_forces: NDArray[np.floating], + external_torques: NDArray[np.floating], + lengths: NDArray[np.floating], + rest_lengths: NDArray[np.floating], + tangents: NDArray[np.floating], + dilatation: NDArray[np.floating], + dilatation_rate: NDArray[np.floating], + voronoi_dilatation: NDArray[np.floating], + rest_voronoi_lengths: NDArray[np.floating], + sigma: NDArray[np.floating], + kappa: NDArray[np.floating], + rest_sigma: NDArray[np.floating], + rest_kappa: NDArray[np.floating], + internal_stress: NDArray[np.floating], + internal_couple: NDArray[np.floating], + ring_rod_flag: bool, ): self.n_elems = n_elements self.position_collection = position @@ -242,17 +246,17 @@ def __init__( def straight_rod( cls, n_elements: int, - start: np.ndarray, - direction: np.ndarray, - normal: np.ndarray, + start: NDArray[np.floating], + direction: NDArray[np.floating], + normal: NDArray[np.floating], base_length: float, base_radius: float, density: float, *, nu: Optional[float] = None, youngs_modulus: float, - **kwargs, - ): + **kwargs: Any, + ) -> Self: """ Cosserat rod constructor for straight-rod geometry. @@ -390,17 +394,17 @@ def straight_rod( def ring_rod( cls, n_elements: int, - ring_center_position: np.ndarray, - direction: np.ndarray, - normal: np.ndarray, + ring_center_position: NDArray[np.floating], + direction: NDArray[np.floating], + normal: NDArray[np.floating], base_length: float, base_radius: float, density: float, *, nu: Optional[float] = None, youngs_modulus: float, - **kwargs, - ): + **kwargs: Any, + ) -> Self: """ Cosserat rod constructor for straight-rod geometry. @@ -536,7 +540,7 @@ def ring_rod( rod.REQUISITE_MODULE.append(Constraints) return rod - def compute_internal_forces_and_torques(self, time): + def compute_internal_forces_and_torques(self, time: float) -> None: """ Compute internal forces and torques. We need to compute internal forces and torques before the acceleration because they are used in interaction. Thus in order to speed up simulation, we will compute internal forces and torques @@ -591,7 +595,7 @@ def compute_internal_forces_and_torques(self, time): ) # Interface to time-stepper mixins (Symplectic, Explicit), which calls this method - def update_accelerations(self, time): + def update_accelerations(self, time: float) -> None: """ Updates the acceleration variables @@ -613,12 +617,12 @@ def update_accelerations(self, time): self.dilatation, ) - def zeroed_out_external_forces_and_torques(self, time): + def zeroed_out_external_forces_and_torques(self, time: np.floating) -> None: _zeroed_out_external_forces_and_torques( self.external_forces, self.external_torques ) - def compute_translational_energy(self): + def compute_translational_energy(self) -> NDArray[np.floating]: """ Compute total translational energy of the rod at the instance. """ @@ -632,7 +636,7 @@ def compute_translational_energy(self): ).sum() ) - def compute_rotational_energy(self): + def compute_rotational_energy(self) -> NDArray[np.floating]: """ Compute total rotational energy of the rod at the instance. """ @@ -642,7 +646,7 @@ def compute_rotational_energy(self): ) return 0.5 * np.einsum("ik,ik->k", self.omega_collection, J_omega_upon_e).sum() - def compute_velocity_center_of_mass(self): + def compute_velocity_center_of_mass(self) -> NDArray[np.floating]: """ Compute velocity center of mass of the rod at the instance. """ @@ -651,7 +655,7 @@ def compute_velocity_center_of_mass(self): return sum_mass_times_velocity / self.mass.sum() - def compute_position_center_of_mass(self): + def compute_position_center_of_mass(self) -> NDArray[np.floating]: """ Compute position center of mass of the rod at the instance. """ @@ -660,7 +664,7 @@ def compute_position_center_of_mass(self): return sum_mass_times_position / self.mass.sum() - def compute_bending_energy(self): + def compute_bending_energy(self) -> NDArray[np.floating]: """ Compute total bending energy of the rod at the instance. """ @@ -676,7 +680,7 @@ def compute_bending_energy(self): ).sum() ) - def compute_shear_energy(self): + def compute_shear_energy(self) -> NDArray[np.floating]: """ Compute total shear energy of the rod at the instance. """ @@ -695,8 +699,12 @@ def compute_shear_energy(self): @numba.njit(cache=True) def _compute_geometry_from_state( - position_collection, volume, lengths, tangents, radius -): + position_collection: NDArray[np.floating], + volume: NDArray[np.floating], + lengths: NDArray[np.floating], + tangents: NDArray[np.floating], + radius: NDArray[np.floating], +) -> None: """ Update given . """ @@ -719,16 +727,16 @@ def _compute_geometry_from_state( @numba.njit(cache=True) def _compute_all_dilatations( - position_collection, - volume, - lengths, - tangents, - radius, - dilatation, - rest_lengths, - rest_voronoi_lengths, - voronoi_dilatation, -): + position_collection: NDArray[np.floating], + volume: NDArray[np.floating], + lengths: NDArray[np.floating], + tangents: NDArray[np.floating], + radius: NDArray[np.floating], + dilatation: NDArray[np.floating], + rest_lengths: NDArray[np.floating], + rest_voronoi_lengths: NDArray[np.floating], + voronoi_dilatation: NDArray[np.floating], +) -> None: """ Update """ @@ -749,8 +757,12 @@ def _compute_all_dilatations( @numba.njit(cache=True) def _compute_dilatation_rate( - position_collection, velocity_collection, lengths, rest_lengths, dilatation_rate -): + position_collection: NDArray[np.floating], + velocity_collection: NDArray[np.floating], + lengths: NDArray[np.floating], + rest_lengths: NDArray[np.floating], + dilatation_rate: NDArray[np.floating], +) -> None: """ Update dilatation_rate given position, velocity, length, and rest_length """ @@ -776,18 +788,18 @@ def _compute_dilatation_rate( @numba.njit(cache=True) def _compute_shear_stretch_strains( - position_collection, - volume, - lengths, - tangents, - radius, - rest_lengths, - rest_voronoi_lengths, - dilatation, - voronoi_dilatation, - director_collection, - sigma, -): + position_collection: NDArray[np.floating], + volume: NDArray[np.floating], + lengths: NDArray[np.floating], + tangents: NDArray[np.floating], + radius: NDArray[np.floating], + rest_lengths: NDArray[np.floating], + rest_voronoi_lengths: NDArray[np.floating], + dilatation: NDArray[np.floating], + voronoi_dilatation: NDArray[np.floating], + director_collection: NDArray[np.floating], + sigma: NDArray[np.floating], +) -> None: """ Update given . """ @@ -811,21 +823,21 @@ def _compute_shear_stretch_strains( @numba.njit(cache=True) def _compute_internal_shear_stretch_stresses_from_model( - position_collection, - volume, - lengths, - tangents, - radius, - rest_lengths, - rest_voronoi_lengths, - dilatation, - voronoi_dilatation, - director_collection, - sigma, - rest_sigma, - shear_matrix, - internal_stress, -): + position_collection: NDArray[np.floating], + volume: NDArray[np.floating], + lengths: NDArray[np.floating], + tangents: NDArray[np.floating], + radius: NDArray[np.floating], + rest_lengths: NDArray[np.floating], + rest_voronoi_lengths: NDArray[np.floating], + dilatation: NDArray[np.floating], + voronoi_dilatation: NDArray[np.floating], + director_collection: NDArray[np.floating], + sigma: NDArray[np.floating], + rest_sigma: NDArray[np.floating], + shear_matrix: NDArray[np.floating], + internal_stress: NDArray[np.floating], +) -> None: """ Update given . @@ -850,7 +862,11 @@ def _compute_internal_shear_stretch_stresses_from_model( @numba.njit(cache=True) -def _compute_bending_twist_strains(director_collection, rest_voronoi_lengths, kappa): +def _compute_bending_twist_strains( + director_collection: NDArray[np.floating], + rest_voronoi_lengths: NDArray[np.floating], + kappa: NDArray[np.floating], +) -> None: """ Update given . """ @@ -864,13 +880,13 @@ def _compute_bending_twist_strains(director_collection, rest_voronoi_lengths, ka @numba.njit(cache=True) def _compute_internal_bending_twist_stresses_from_model( - director_collection, - rest_voronoi_lengths, - internal_couple, - bend_matrix, - kappa, - rest_kappa, -): + director_collection: NDArray[np.floating], + rest_voronoi_lengths: NDArray[np.floating], + internal_couple: NDArray[np.floating], + bend_matrix: NDArray[np.floating], + kappa: NDArray[np.floating], + rest_kappa: NDArray[np.floating], +) -> None: """ Upate given . @@ -893,23 +909,23 @@ def _compute_internal_bending_twist_stresses_from_model( @numba.njit(cache=True) def _compute_internal_forces( - position_collection, - volume, - lengths, - tangents, - radius, - rest_lengths, - rest_voronoi_lengths, - dilatation, - voronoi_dilatation, - director_collection, - sigma, - rest_sigma, - shear_matrix, - internal_stress, - internal_forces, - ghost_elems_idx, -): + position_collection: NDArray[np.floating], + volume: NDArray[np.floating], + lengths: NDArray[np.floating], + tangents: NDArray[np.floating], + radius: NDArray[np.floating], + rest_lengths: NDArray[np.floating], + rest_voronoi_lengths: NDArray[np.floating], + dilatation: NDArray[np.floating], + voronoi_dilatation: NDArray[np.floating], + director_collection: NDArray[np.floating], + sigma: NDArray[np.floating], + rest_sigma: NDArray[np.floating], + shear_matrix: NDArray[np.floating], + internal_stress: NDArray[np.floating], + internal_forces: NDArray[np.floating], + ghost_elems_idx: NDArray[np.floating], +) -> None: """ Update given . """ @@ -954,26 +970,26 @@ def _compute_internal_forces( @numba.njit(cache=True) def _compute_internal_torques( - position_collection, - velocity_collection, - tangents, - lengths, - rest_lengths, - director_collection, - rest_voronoi_lengths, - bend_matrix, - rest_kappa, - kappa, - voronoi_dilatation, - mass_second_moment_of_inertia, - omega_collection, - internal_stress, - internal_couple, - dilatation, - dilatation_rate, - internal_torques, - ghost_voronoi_idx, -): + position_collection: NDArray[np.floating], + velocity_collection: NDArray[np.floating], + tangents: NDArray[np.floating], + lengths: NDArray[np.floating], + rest_lengths: NDArray[np.floating], + director_collection: NDArray[np.floating], + rest_voronoi_lengths: NDArray[np.floating], + bend_matrix: NDArray[np.floating], + rest_kappa: NDArray[np.floating], + kappa: NDArray[np.floating], + voronoi_dilatation: NDArray[np.floating], + mass_second_moment_of_inertia: NDArray[np.floating], + omega_collection: NDArray[np.floating], + internal_stress: NDArray[np.floating], + internal_couple: NDArray[np.floating], + dilatation: NDArray[np.floating], + dilatation_rate: NDArray[np.floating], + internal_torques: NDArray[np.floating], + ghost_voronoi_idx: NDArray[np.integer], +) -> None: """ Update . """ @@ -1043,16 +1059,16 @@ def _compute_internal_torques( @numba.njit(cache=True) def _update_accelerations( - acceleration_collection, - internal_forces, - external_forces, - mass, - alpha_collection, - inv_mass_second_moment_of_inertia, - internal_torques, - external_torques, - dilatation, -): + acceleration_collection: NDArray[np.floating], + internal_forces: NDArray[np.floating], + external_forces: NDArray[np.floating], + mass: NDArray[np.floating], + alpha_collection: NDArray[np.floating], + inv_mass_second_moment_of_inertia: NDArray[np.floating], + internal_torques: NDArray[np.floating], + external_torques: NDArray[np.floating], + dilatation: NDArray[np.floating], +) -> None: """ Update given . """ @@ -1077,7 +1093,9 @@ def _update_accelerations( @numba.njit(cache=True) -def _zeroed_out_external_forces_and_torques(external_forces, external_torques): +def _zeroed_out_external_forces_and_torques( + external_forces: NDArray[np.floating], external_torques: NDArray[np.floating] +) -> None: """ This function is to zeroed out external forces and torques. diff --git a/elastica/rod/data_structures.py b/elastica/rod/data_structures.py index 3c8b40328..c65a9c4b5 100644 --- a/elastica/rod/data_structures.py +++ b/elastica/rod/data_structures.py @@ -1,6 +1,9 @@ __doc__ = "Data structure wrapper for rod components" +from typing import Any, Optional +from typing_extensions import Self import numpy as np +from numpy.typing import NDArray from numba import njit from elastica._rotations import _get_rotation_matrix, _rotate from elastica._linalg import _batch_matmul @@ -9,7 +12,7 @@ # FIXME : Explicit Stepper doesn't work as States lose the # views they initially had when working with a timestepper. # class _RodExplicitStepperMixin: -# def __init__(self): +# def __init__(self) -> None: # ( # self.state, # self.__deriv_state, @@ -43,7 +46,7 @@ class _RodSymplecticStepperMixin: - def __init__(self): + def __init__(self) -> None: self.kinematic_states = _KinematicState( self.position_collection, self.director_collection ) @@ -60,18 +63,40 @@ def __init__(self): # is another function self.kinematic_rates = self.dynamic_states.kinematic_rates - def update_internal_forces_and_torques(self, time, *args, **kwargs): + def update_internal_forces_and_torques( + self, time: np.floating, *args: Any, **kwargs: Any + ) -> None: self.compute_internal_forces_and_torques(time) - def dynamic_rates(self, time, prefac, *args, **kwargs): + def dynamic_rates( + self, time: np.floating, prefac: np.floating, *args: Any, **kwargs: Any + ) -> NDArray[np.floating]: self.update_accelerations(time) return self.dynamic_states.dynamic_rates(time, prefac, *args, **kwargs) - def reset_external_forces_and_torques(self, time, *args, **kwargs): + def reset_external_forces_and_torques( + self, time: np.floating, *args: Any, **kwargs: Any + ) -> None: self.zeroed_out_external_forces_and_torques(time) -def _bootstrap_from_data(stepper_type: str, n_elems: int, vector_states, matrix_states): +def _bootstrap_from_data( + stepper_type: str, + n_elems: int, + vector_states: NDArray[np.floating], + matrix_states: NDArray[np.floating], +) -> Optional[ + tuple[ + "_State", + "_DerivativeState", + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + ] +]: """Returns states wrapping numpy arrays based on the time-stepping algorithm Convenience method that takes in rod internal (raw np.ndarray) data, create views @@ -119,7 +144,7 @@ def _bootstrap_from_data(stepper_type: str, n_elems: int, vector_states, matrix_ # ) raise NotImplementedError else: - return + return None n_velocity_end = n_nodes + n_nodes velocity = np.ndarray.view(vector_states[..., n_nodes:n_velocity_end]) @@ -154,10 +179,10 @@ class _State: def __init__( self, n_elems: int, - position_collection_view, - director_collection_view, - kinematic_rate_collection_view, - ): + position_collection_view: NDArray[np.floating], + director_collection_view: NDArray[np.floating], + kinematic_rate_collection_view: NDArray[np.floating], + ) -> None: """ Parameters ---------- @@ -173,7 +198,7 @@ def __init__( self.director_collection = director_collection_view self.kinematic_rate_collection = kinematic_rate_collection_view - def __iadd__(self, scaled_deriv_array): + def __iadd__(self, scaled_deriv_array: NDArray[np.floating]) -> Self: """overloaded += operator The add for directors is customized to reflect Rodrigues' rotation @@ -242,7 +267,7 @@ def __iadd__(self, scaled_deriv_array): ] return self - def __add__(self, scaled_derivative_state): + def __add__(self, scaled_derivative_state: NDArray[np.floating]) -> "_State": """overloaded + operator, useful in state.k1 = state + dt * deriv_state The add for directors is customized to reflect Rodrigues' rotation @@ -296,7 +321,9 @@ class _DerivativeState: /multiplication used. """ - def __init__(self, _unused_n_elems: int, rate_collection_view): + def __init__( + self, _unused_n_elems: int, rate_collection_view: NDArray[np.floating] + ) -> None: """ Parameters ---------- @@ -307,7 +334,7 @@ def __init__(self, _unused_n_elems: int, rate_collection_view): super(_DerivativeState, self).__init__() self.rate_collection = rate_collection_view - def __rmul__(self, scalar): + def __rmul__(self, scalar: np.floating) -> NDArray[np.floating]: """overloaded scalar * self, Parameters @@ -355,7 +382,7 @@ def __rmul__(self, scalar): """ return scalar * self.rate_collection - def __mul__(self, scalar): + def __mul__(self, scalar: np.floating) -> NDArray[np.floating]: """overloaded self * scalar TODO Check if this pattern (forwarding to __mul__) has @@ -388,7 +415,11 @@ class _KinematicState: only these methods are provided. """ - def __init__(self, position_collection_view, director_collection_view): + def __init__( + self, + position_collection_view: NDArray[np.floating], + director_collection_view: NDArray[np.floating], + ) -> None: """ Parameters ---------- @@ -403,13 +434,13 @@ def __init__(self, position_collection_view, director_collection_view): @njit(cache=True) def overload_operator_kinematic_numba( - n_nodes, - prefac, - position_collection, - director_collection, - velocity_collection, - omega_collection, -): + n_nodes: int, + prefac: np.floating, + position_collection: NDArray[np.floating], + director_collection: NDArray[np.floating], + velocity_collection: NDArray[np.floating], + omega_collection: NDArray[np.floating], +) -> None: """overloaded += operator The add for directors is customized to reflect Rodrigues' rotation @@ -449,11 +480,11 @@ class _DynamicState: def __init__( self, - v_w_collection, - dvdt_dwdt_collection, - velocity_collection, - omega_collection, - ): + v_w_collection: NDArray[np.floating], + dvdt_dwdt_collection: NDArray[np.floating], + velocity_collection: NDArray[np.floating], + omega_collection: NDArray[np.floating], + ) -> None: """ Parameters ---------- @@ -470,7 +501,9 @@ def __init__( self.velocity_collection = velocity_collection self.omega_collection = omega_collection - def kinematic_rates(self, time, prefac): + def kinematic_rates( + self, time: np.floating, prefac: np.floating + ) -> tuple[NDArray[np.floating], NDArray[np.floating]]: """Yields kinematic rates to interact with _KinematicState Returns @@ -486,7 +519,9 @@ def kinematic_rates(self, time, prefac): # Comes from kin_state -> (x,Q) += dt * (v,w) <- First part of dyn_state return self.velocity_collection, self.omega_collection - def dynamic_rates(self, time, prefac): + def dynamic_rates( + self, time: np.floating, prefac: np.floating + ) -> NDArray[np.floating]: """Yields dynamic rates to add to with _DynamicState Returns ------- @@ -501,7 +536,10 @@ def dynamic_rates(self, time, prefac): @njit(cache=True) -def overload_operator_dynamic_numba(rate_collection, scaled_second_deriv_array): +def overload_operator_dynamic_numba( + rate_collection: NDArray[np.floating], + scaled_second_deriv_array: NDArray[np.floating], +) -> None: """overloaded += operator, updating dynamic_rates Parameters ---------- diff --git a/elastica/rod/factory_function.py b/elastica/rod/factory_function.py index 6eca6b5ac..22cc33bc6 100644 --- a/elastica/rod/factory_function.py +++ b/elastica/rod/factory_function.py @@ -1,30 +1,64 @@ __doc__ = """ Factory function to allocate variables for Cosserat Rod""" -from typing import Optional, Tuple +from typing import Any, Optional, Tuple import logging import numpy as np from numpy.testing import assert_allclose +from numpy.typing import NDArray from elastica.utils import MaxDimension, Tolerance from elastica._linalg import _batch_cross, _batch_norm, _batch_dot def allocate( - n_elements, - direction, - normal, - base_length, - base_radius, - density, - youngs_modulus: float, + n_elements: int, + direction: NDArray[np.floating], + normal: NDArray[np.floating], + base_length: np.floating, + base_radius: np.floating, + density: np.floating, + youngs_modulus: np.floating, *, rod_origin_position: np.ndarray, ring_rod_flag: bool, - shear_modulus: Optional[float] = None, + shear_modulus: Optional[np.floating] = None, position: Optional[np.ndarray] = None, directors: Optional[np.ndarray] = None, rest_sigma: Optional[np.ndarray] = None, rest_kappa: Optional[np.ndarray] = None, - **kwargs, -): + **kwargs: Any, +) -> tuple[ + int, + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], + NDArray[np.floating], +]: log = logging.getLogger() if "poisson_ratio" in kwargs: @@ -335,14 +369,16 @@ def allocate( """ -def _assert_dim(vector, max_dim: int, name: str): +def _assert_dim(vector: np.ndarray, max_dim: int, name: str) -> None: assert vector.ndim < max_dim, ( f"Input {name} dimension is not correct {vector.shape}" + f" It should be maximum {max_dim}D vector or single floating number." ) -def _assert_shape(array: np.ndarray, expected_shape: Tuple[int], name: str): +def _assert_shape( + array: np.ndarray, expected_shape: Tuple[int, ...], name: str +) -> None: assert array.shape == expected_shape, ( f"Given {name} shape is not correct, it should be " + str(expected_shape) @@ -351,7 +387,9 @@ def _assert_shape(array: np.ndarray, expected_shape: Tuple[int], name: str): ) -def _position_validity_checker(position, start, n_elements): +def _position_validity_checker( + position: NDArray[np.floating], start: NDArray[np.floating], n_elements: int +) -> None: """Checker on user-defined position validity""" _assert_shape(position, (MaxDimension.value(), n_elements + 1), "position") @@ -367,7 +405,9 @@ def _position_validity_checker(position, start, n_elements): ) -def _directors_validity_checker(directors, tangents, n_elements): +def _directors_validity_checker( + directors: NDArray[np.floating], tangents: NDArray[np.floating], n_elements: int +) -> None: """Checker on user-defined directors validity""" _assert_shape( directors, (MaxDimension.value(), MaxDimension.value(), n_elements), "directors" @@ -413,7 +453,11 @@ def _directors_validity_checker(directors, tangents, n_elements): ) -def _position_validity_checker_ring_rod(position, ring_center_position, n_elements): +def _position_validity_checker_ring_rod( + position: NDArray[np.floating], + ring_center_position: NDArray[np.floating], + n_elements: int, +) -> None: """Checker on user-defined position validity""" _assert_shape(position, (MaxDimension.value(), n_elements), "position") diff --git a/elastica/rod/knot_theory.py b/elastica/rod/knot_theory.py index 23d2b009f..537d8d4e5 100644 --- a/elastica/rod/knot_theory.py +++ b/elastica/rod/knot_theory.py @@ -14,6 +14,7 @@ from numba import njit import numpy as np +from numpy.typing import NDArray from elastica.rod.rod_base import RodBase from elastica._linalg import _batch_norm, _batch_dot, _batch_cross @@ -38,7 +39,7 @@ def radius(self) -> np.ndarray: ... def base_length(self) -> np.ndarray: ... -class KnotTheory: +class KnotTheory(KnotTheoryCompatibleProtocol): """ This mixin should be used in RodBase-derived class that satisfies KnotCompatibleProtocol. The theory behind this module is based on the method from Klenin & Langowski 2000 paper. @@ -46,7 +47,7 @@ class KnotTheory: KnotTheory can be mixed with any rod-class based on RodBase:: class MyRod(RodBase, KnotTheory): - def __init__(self): + def __init__(self) -> None: super().__init__() rod = MyRod(...) @@ -78,7 +79,7 @@ def __init__(self): MIXIN_PROTOCOL = Union[RodBase, KnotTheoryCompatibleProtocol] - def compute_twist(self: MIXIN_PROTOCOL): + def compute_twist(self: MIXIN_PROTOCOL) -> NDArray[np.floating]: """ See :ref:`api/rods:Knot Theory (Mixin)` for the detail. """ @@ -91,8 +92,8 @@ def compute_twist(self: MIXIN_PROTOCOL): def compute_writhe( self: MIXIN_PROTOCOL, type_of_additional_segment: str = "next_tangent", - alpha: float = 1.0, - ): + alpha: np.floating = 1.0, + ) -> NDArray[np.floating]: """ See :ref:`api/rods:Knot Theory (Mixin)` for the detail. @@ -114,8 +115,8 @@ def compute_writhe( def compute_link( self: MIXIN_PROTOCOL, type_of_additional_segment: str = "next_tangent", - alpha: float = 1.0, - ): + alpha: np.floating = 1.0, + ) -> NDArray[np.floating]: """ See :ref:`api/rods:Knot Theory (Mixin)` for the detail. @@ -138,7 +139,9 @@ def compute_link( )[0] -def compute_twist(center_line, normal_collection): +def compute_twist( + center_line: NDArray[np.floating], normal_collection: NDArray[np.floating] +) -> tuple[NDArray[np.floating], NDArray[np.floating]]: """ Compute the twist of a rod, using center_line and normal collection. @@ -189,7 +192,9 @@ def compute_twist(center_line, normal_collection): @njit(cache=True) -def _compute_twist(center_line, normal_collection): +def _compute_twist( + center_line: NDArray[np.floating], normal_collection: NDArray[np.floating] +) -> tuple[NDArray[np.floating], NDArray[np.floating]]: """ Parameters ---------- @@ -264,7 +269,11 @@ def _compute_twist(center_line, normal_collection): return total_twist, local_twist -def compute_writhe(center_line, segment_length, type_of_additional_segment): +def compute_writhe( + center_line: NDArray[np.floating], + segment_length: np.floating, + type_of_additional_segment: str, +) -> NDArray[np.floating]: """ This function computes the total writhe history of a rod. @@ -314,7 +323,7 @@ def compute_writhe(center_line, segment_length, type_of_additional_segment): @njit(cache=True) -def _compute_writhe(center_line): +def _compute_writhe(center_line: NDArray[np.floating]) -> NDArray[np.floating]: """ Parameters ---------- @@ -386,12 +395,12 @@ def _compute_writhe(center_line): def compute_link( - center_line: np.ndarray, - normal_collection: np.ndarray, - radius: np.ndarray, - segment_length: float, + center_line: NDArray[np.floating], + normal_collection: NDArray[np.floating], + radius: NDArray[np.floating], + segment_length: np.floating, type_of_additional_segment: str, -): +) -> NDArray[np.floating]: """ This function computes the total link history of a rod. @@ -470,7 +479,11 @@ def compute_link( @njit(cache=True) -def _compute_auxiliary_line(center_line, normal_collection, radius): +def _compute_auxiliary_line( + center_line: NDArray[np.floating], + normal_collection: NDArray[np.floating], + radius: NDArray[np.floating], +) -> NDArray[np.floating]: """ This function computes the auxiliary line using rod center line and normal collection. @@ -525,7 +538,9 @@ def _compute_auxiliary_line(center_line, normal_collection, radius): @njit(cache=True) -def _compute_link(center_line, auxiliary_line): +def _compute_link( + center_line: NDArray[np.floating], auxiliary_line: NDArray[np.floating] +) -> NDArray[np.floating]: """ Parameters @@ -604,8 +619,11 @@ def _compute_link(center_line, auxiliary_line): @njit(cache=True) def _compute_auxiliary_line_added_segments( - beginning_direction, end_direction, auxiliary_line, segment_length -): + beginning_direction: NDArray[np.floating], + end_direction: NDArray[np.floating], + auxiliary_line: NDArray[np.floating], + segment_length: np.floating, +) -> NDArray[np.floating]: """ This code is for computing position of added segments to the auxiliary line. @@ -647,8 +665,10 @@ def _compute_auxiliary_line_added_segments( @njit(cache=True) def _compute_additional_segment( - center_line, segment_length, type_of_additional_segment -): + center_line: NDArray[np.floating], + segment_length: np.floating, + type_of_additional_segment: str, +) -> tuple[NDArray[np.floating], NDArray[np.floating], NDArray[np.floating]]: """ This function adds two points at the end of center line. Distance from the center line is given by segment_length. Direction from center line to the new point locations can be computed using 3 methods, which can be selected by diff --git a/elastica/rod/rod_base.py b/elastica/rod/rod_base.py index 314efe6dc..d49c0431f 100644 --- a/elastica/rod/rod_base.py +++ b/elastica/rod/rod_base.py @@ -1,5 +1,9 @@ __doc__ = """Base class for rods""" +from typing import Any +import numpy as np +from numpy.typing import NDArray + class RodBase: """ @@ -11,10 +15,15 @@ class RodBase: """ - REQUISITE_MODULES = [] + REQUISITE_MODULES: list[Any] = [] - def __init__(self): + def __init__(self) -> None: """ RodBase does not take any arguments. """ - pass + self.position_collection: NDArray[np.floating] + self.omega_collection: NDArray[np.floating] + self.acceleration_collection: NDArray[np.floating] + self.alpha_collection: NDArray[np.floating] + self.external_forces: NDArray[np.floating] + self.external_torques: NDArray[np.floating] From 856e901f8a63aca3b0a45b4d89936de77115f0a2 Mon Sep 17 00:00:00 2001 From: Ankith Date: Mon, 20 May 2024 11:51:55 +0530 Subject: [PATCH 054/134] KnotTheory remove protocol inherit Co-authored-by: Seung Hyun Kim --- elastica/rod/knot_theory.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/elastica/rod/knot_theory.py b/elastica/rod/knot_theory.py index 537d8d4e5..ef469257d 100644 --- a/elastica/rod/knot_theory.py +++ b/elastica/rod/knot_theory.py @@ -39,7 +39,7 @@ def radius(self) -> np.ndarray: ... def base_length(self) -> np.ndarray: ... -class KnotTheory(KnotTheoryCompatibleProtocol): +class KnotTheory: """ This mixin should be used in RodBase-derived class that satisfies KnotCompatibleProtocol. The theory behind this module is based on the method from Klenin & Langowski 2000 paper. From 0fd9c1c4d37d1de5914aa109c0520b3dabe65bbd Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Mon, 20 May 2024 15:15:13 -0500 Subject: [PATCH 055/134] test: fix missing values --- elastica/modules/base_system.py | 2 +- elastica/modules/connections.py | 8 ++++---- elastica/modules/protocol.py | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index ed1b95cfd..1bf549d1d 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -162,7 +162,7 @@ def _get_sys_idx_if_valid(self, sys_to_be_added: SystemType) -> SystemIdxType: @final def blocks(self) -> Generator[SystemType, None, None]: - assert self._finalize_flag, "The simulator is not finalized." + # assert self._finalize_flag, "The simulator is not finalized." for block in self._memory_blocks: yield block diff --git a/elastica/modules/connections.py b/elastica/modules/connections.py index 6576d9688..92966b782 100644 --- a/elastica/modules/connections.py +++ b/elastica/modules/connections.py @@ -178,9 +178,9 @@ def __init__( self._second_sys_idx: SystemIdxType = second_sys_idx self._first_sys_n_lim: int = first_sys_nlim self._second_sys_n_lim: int = second_sys_nlim - self._connect_cls: Type[FreeJoint] - self.first_sys_connection_idx: ConnectionIndex - self.second_sys_connection_idx: ConnectionIndex + self._connect_cls: Type[FreeJoint] = None + self.first_sys_connection_idx: ConnectionIndex = None + self.second_sys_connection_idx: ConnectionIndex = None def set_index( self, first_idx: ConnectionIndex, second_idx: ConnectionIndex @@ -298,7 +298,7 @@ def id( ) def instantiate(self) -> FreeJoint: - if not self._connect_cls: + if self._connect_cls is None: raise RuntimeError( "No connections provided to link rod id {0}" "(at {2}) and {1} (at {3}), but a Connection" diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index 424423d0a..9de78e0a2 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -1,4 +1,4 @@ -from typing import Protocol, Generator +from typing import Protocol, Generator, TypeVar, Any from typing_extensions import Self # 3.11: from typing import Self from elastica.typing import ( SystemIdxType, @@ -34,7 +34,7 @@ def blocks(self) -> Generator[SystemType, None, None]: ... def _get_sys_idx_if_valid(self, sys_to_be_added: SystemType) -> SystemIdxType: ... -M = TypeVar("M", bound=ModuleProtocol) +M = TypeVar("M", bound="ModuleProtocol") class ModuleProtocol(Protocol[M]): From 424a9a8bac66bd640b3b5bdd7283c47f8800e000 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Mon, 20 May 2024 16:06:57 -0500 Subject: [PATCH 056/134] ignore njit wrapper Exactly tracing the type for njit dispatched function will only make the code unreadable. --- elastica/_calculus.py | 25 +++++++++--------- elastica/_contact_functions.py | 17 ++++++------ elastica/_linalg.py | 22 ++++++++-------- elastica/_rotations.py | 6 ++--- elastica/_synchronize_periodic_boundary.py | 26 ++++++++++--------- elastica/boundary_conditions.py | 18 ++++++------- elastica/dissipation.py | 6 ++--- .../parallel_connection.py | 4 +-- elastica/experimental/interaction.py | 2 +- elastica/external_forces.py | 10 +++---- elastica/interaction.py | 8 +++--- .../_reset_ghost_vector_or_scalar.py | 4 +-- elastica/rod/data_structures.py | 4 +-- elastica/rod/knot_theory.py | 12 ++++----- 14 files changed, 84 insertions(+), 80 deletions(-) diff --git a/elastica/_calculus.py b/elastica/_calculus.py index ae730da15..edd7afdee 100644 --- a/elastica/_calculus.py +++ b/elastica/_calculus.py @@ -1,9 +1,10 @@ __doc__ = """ Quadrature and difference kernels """ -from typing import Any, Union +from typing import Union import numpy as np from numpy import zeros, empty from numpy.typing import NDArray -from numba import njit +import numba +from numba import njit, float64 from elastica.reset_functions_for_block_structure._reset_ghost_vector_or_scalar import ( _reset_vector_ghost, ) @@ -20,8 +21,8 @@ def _get_zero_array(dim: int, ndim: int) -> Union[float, NDArray[np.floating], N return None -@njit(cache=True) -def _trapezoidal(array_collection: NDArray[np.floating]) -> NDArray[np.floating]: +@njit(cache=True) # type: ignore +def _trapezoidal(array_collection: NDArray[np.float64]) -> NDArray[np.float64]: """ Simple trapezoidal quadrature rule with zero at end-points, in a dimension agnostic way @@ -66,7 +67,7 @@ def _trapezoidal(array_collection: NDArray[np.floating]) -> NDArray[np.floating] return temp_collection -@njit(cache=True) +@njit(cache=True) # type: ignore def _trapezoidal_for_block_structure( array_collection: NDArray[np.floating], ghost_idx: NDArray[np.integer] ) -> NDArray[np.floating]: @@ -120,7 +121,7 @@ def _trapezoidal_for_block_structure( return temp_collection -@njit(cache=True) +@njit(cache=True) # type: ignore def _two_point_difference( array_collection: NDArray[np.floating], ) -> NDArray[np.floating]: @@ -163,7 +164,7 @@ def _two_point_difference( return temp_collection -@njit(cache=True) +@njit(cache=True) # type: ignore def _two_point_difference_for_block_structure( array_collection: NDArray[np.floating], ghost_idx: NDArray[np.integer] ) -> NDArray[np.floating]: @@ -216,7 +217,7 @@ def _two_point_difference_for_block_structure( return temp_collection -@njit(cache=True) +@njit(cache=True) # type: ignore def _difference(vector: NDArray[np.floating]) -> NDArray[np.floating]: """ This function computes difference between elements of a batch vector. @@ -247,7 +248,7 @@ def _difference(vector: NDArray[np.floating]) -> NDArray[np.floating]: return output_vector -@njit(cache=True) +@njit(cache=True) # type: ignore def _average(vector: NDArray[np.floating]) -> NDArray[np.floating]: """ This function computes the average between elements of a vector. @@ -277,7 +278,7 @@ def _average(vector: NDArray[np.floating]) -> NDArray[np.floating]: return output_vector -@njit(cache=True) +@njit(cache=True) # type: ignore def _clip_array( input_array: NDArray[np.floating], vmin: np.floating, vmax: np.floating ) -> NDArray[np.floating]: @@ -315,8 +316,8 @@ def _clip_array( return input_array -@njit(cache=True) -def _isnan_check(array: NDArray[Any]) -> bool: +@njit(cache=True) # type: ignore +def _isnan_check(array: NDArray) -> bool: """ This function checks if there is any nan inside the array. If there is nan, it returns True boolean. diff --git a/elastica/_contact_functions.py b/elastica/_contact_functions.py index 3b2d25d5a..be33436ac 100644 --- a/elastica/_contact_functions.py +++ b/elastica/_contact_functions.py @@ -22,12 +22,13 @@ _batch_matrix_transpose, _batch_vec_oneD_vec_cross, ) -import numba import numpy as np from numpy.typing import NDArray +from numba import njit -@numba.njit(cache=True) + +@njit(cache=True) # type: ignore def _calculate_contact_forces_rod_cylinder( x_collection_rod: NDArray[np.floating], edge_collection_rod: NDArray[np.floating], @@ -154,7 +155,7 @@ def _calculate_contact_forces_rod_cylinder( ) -@numba.njit(cache=True) +@njit(cache=True) # type: ignore def _calculate_contact_forces_rod_rod( x_collection_rod_one: NDArray[np.floating], radius_rod_one: NDArray[np.floating], @@ -271,7 +272,7 @@ def _calculate_contact_forces_rod_rod( external_forces_rod_two[..., j + 1] += net_contact_force -@numba.njit(cache=True) +@njit(cache=True) # type: ignore def _calculate_contact_forces_self_rod( x_collection_rod: NDArray[np.floating], radius_rod: NDArray[np.floating], @@ -359,7 +360,7 @@ def _calculate_contact_forces_self_rod( external_forces_rod[..., j + 1] += net_contact_force -@numba.njit(cache=True) +@njit(cache=True) # type: ignore def _calculate_contact_forces_rod_sphere( x_collection_rod: NDArray[np.floating], edge_collection_rod: NDArray[np.floating], @@ -485,7 +486,7 @@ def _calculate_contact_forces_rod_sphere( ) -@numba.njit(cache=True) +@njit(cache=True) # type: ignore def _calculate_contact_forces_rod_plane( plane_origin: NDArray[np.floating], plane_normal: NDArray[np.floating], @@ -570,7 +571,7 @@ def _calculate_contact_forces_rod_plane( return (_batch_norm(plane_response_force), no_contact_point_idx) -@numba.njit(cache=True) +@njit(cache=True) # type: ignore def _calculate_contact_forces_rod_plane_with_anisotropic_friction( plane_origin: NDArray[np.floating], plane_normal: NDArray[np.floating], @@ -783,7 +784,7 @@ def _calculate_contact_forces_rod_plane_with_anisotropic_friction( ) -@numba.njit(cache=True) +@njit(cache=True) # type: ignore def _calculate_contact_forces_cylinder_plane( plane_origin: NDArray[np.floating], plane_normal: NDArray[np.floating], diff --git a/elastica/_linalg.py b/elastica/_linalg.py index 3123ff75b..5097c7093 100644 --- a/elastica/_linalg.py +++ b/elastica/_linalg.py @@ -28,7 +28,7 @@ def levi_civita_tensor(dim: int) -> NDArray[np.floating]: return epsilon -@njit(cache=True) +@njit(cache=True) # type: ignore def _batch_matvec( matrix_collection: NDArray[np.floating], vector_collection: NDArray[np.floating] ) -> NDArray[np.floating]: @@ -61,7 +61,7 @@ def _batch_matvec( return output_vector -@njit(cache=True) +@njit(cache=True) # type: ignore def _batch_matmul( first_matrix_collection: NDArray[np.floating], second_matrix_collection: NDArray[np.floating], @@ -98,7 +98,7 @@ def _batch_matmul( return output_matrix -@njit(cache=True) +@njit(cache=True) # type: ignore def _batch_cross( first_vector_collection: NDArray[np.floating], second_vector_collection: NDArray[np.floating], @@ -141,7 +141,7 @@ def _batch_cross( return output_vector -@njit(cache=True) +@njit(cache=True) # type: ignore def _batch_vec_oneD_vec_cross( first_vector_collection: NDArray[np.floating], second_vector: NDArray[np.floating] ) -> NDArray[np.floating]: @@ -187,7 +187,7 @@ def _batch_vec_oneD_vec_cross( return output_vector -@njit(cache=True) +@njit(cache=True) # type: ignore def _batch_dot( first_vector: NDArray[np.floating], second_vector: NDArray[np.floating] ) -> NDArray[np.floating]: @@ -216,7 +216,7 @@ def _batch_dot( return output_vector -@njit(cache=True) +@njit(cache=True) # type: ignore def _batch_norm(vector: NDArray[np.floating]) -> NDArray[np.floating]: """ This function computes norm of a batch vector @@ -245,7 +245,7 @@ def _batch_norm(vector: NDArray[np.floating]) -> NDArray[np.floating]: return output_vector -@njit(cache=True) +@njit(cache=True) # type: ignore def _batch_product_i_k_to_ik( vector1: NDArray[np.floating], vector2: NDArray[np.floating] ) -> NDArray[np.floating]: @@ -276,7 +276,7 @@ def _batch_product_i_k_to_ik( return output_vector -@njit(cache=True) +@njit(cache=True) # type: ignore def _batch_product_i_ik_to_k( vector1: NDArray[np.floating], vector2: NDArray[np.floating] ) -> NDArray[np.floating]: @@ -309,7 +309,7 @@ def _batch_product_i_ik_to_k( return output_vector -@njit(cache=True) +@njit(cache=True) # type: ignore def _batch_product_k_ik_to_ik( vector1: NDArray[np.floating], vector2: NDArray[np.floating] ) -> NDArray[np.floating]: @@ -340,7 +340,7 @@ def _batch_product_k_ik_to_ik( return output_vector -@njit(cache=True) +@njit(cache=True) # type: ignore def _batch_vector_sum( vector1: NDArray[np.floating], vector2: NDArray[np.floating] ) -> NDArray[np.floating]: @@ -372,7 +372,7 @@ def _batch_vector_sum( return output_vector -@njit(cache=True) +@njit(cache=True) # type: ignore def _batch_matrix_transpose(input_matrix: NDArray[np.floating]) -> NDArray[np.floating]: """ This function takes an batch input matrix and transpose it. diff --git a/elastica/_rotations.py b/elastica/_rotations.py index 7fb9d0633..b7d1180a6 100644 --- a/elastica/_rotations.py +++ b/elastica/_rotations.py @@ -15,7 +15,7 @@ from elastica._linalg import _batch_matmul -@njit(cache=True) +@njit(cache=True) # type: ignore def _get_rotation_matrix( scale: np.floating, axis_collection: NDArray[np.floating] ) -> NDArray[np.floating]: @@ -51,7 +51,7 @@ def _get_rotation_matrix( return rot_mat -@njit(cache=True) +@njit(cache=True) # type: ignore def _rotate( director_collection: NDArray[np.floating], scale: np.floating, @@ -80,7 +80,7 @@ def _rotate( ) -@njit(cache=True) +@njit(cache=True) # type: ignore def _inv_rotate(director_collection: NDArray[np.floating]) -> NDArray[np.floating]: """ Calculated rate of change using Rodrigues' formula diff --git a/elastica/_synchronize_periodic_boundary.py b/elastica/_synchronize_periodic_boundary.py index 0a06622c8..7d4c4882b 100644 --- a/elastica/_synchronize_periodic_boundary.py +++ b/elastica/_synchronize_periodic_boundary.py @@ -10,15 +10,15 @@ from elastica.typing import SystemType -@njit(cache=True) +@njit(cache=True) # type: ignore def _synchronize_periodic_boundary_of_vector_collection( - input: NDArray[np.floating], periodic_idx: NDArray[np.floating] + input_array: NDArray[np.floating], periodic_idx: NDArray[np.floating] ) -> None: """ This function synchronizes the periodic boundaries of a vector collection. Parameters ---------- - input : numpy.ndarray + input_array : numpy.ndarray 2D (dim, blocksize) array containing data with 'float' type. Vector that is going to be synched. periodic_idx : numpy.ndarray 2D (2, n_periodic_boundary) array containing data with 'float' type. Vector containing periodic boundary @@ -30,18 +30,18 @@ def _synchronize_periodic_boundary_of_vector_collection( """ for i in range(3): for k in range(periodic_idx.shape[1]): - input[i, periodic_idx[0, k]] = input[i, periodic_idx[1, k]] + input_array[i, periodic_idx[0, k]] = input_array[i, periodic_idx[1, k]] -@njit(cache=True) +@njit(cache=True) # type: ignore def _synchronize_periodic_boundary_of_matrix_collection( - input: NDArray[np.floating], periodic_idx: NDArray[np.floating] + input_array: NDArray[np.floating], periodic_idx: NDArray[np.floating] ) -> None: """ This function synchronizes the periodic boundaries of a matrix collection. Parameters ---------- - input : numpy.ndarray + input_array : numpy.ndarray 2D (dim, dim, blocksize) array containing data with 'float' type. Matrix collection that is going to be synched. periodic_idx : numpy.ndarray 2D (2, n_periodic_boundary) array containing data with 'float' type. Vector containing periodic boundary @@ -54,19 +54,21 @@ def _synchronize_periodic_boundary_of_matrix_collection( for i in range(3): for j in range(3): for k in range(periodic_idx.shape[1]): - input[i, j, periodic_idx[0, k]] = input[i, j, periodic_idx[1, k]] + input_array[i, j, periodic_idx[0, k]] = input_array[ + i, j, periodic_idx[1, k] + ] -@njit(cache=True) +@njit(cache=True) # type: ignore def _synchronize_periodic_boundary_of_scalar_collection( - input: NDArray[np.floating], periodic_idx: NDArray[np.floating] + input_array: NDArray[np.floating], periodic_idx: NDArray[np.floating] ) -> None: """ This function synchronizes the periodic boundaries of a scalar collection. Parameters ---------- - input : numpy.ndarray + input_array : numpy.ndarray 2D (dim, dim, blocksize) array containing data with 'float' type. Scalar collection that is going to be synched. periodic_idx : numpy.ndarray 2D (2, n_periodic_boundary) array containing data with 'float' type. Vector containing periodic boundary @@ -77,7 +79,7 @@ def _synchronize_periodic_boundary_of_scalar_collection( """ for k in range(periodic_idx.shape[1]): - input[periodic_idx[0, k]] = input[periodic_idx[1, k]] + input_array[periodic_idx[0, k]] = input_array[periodic_idx[1, k]] class _ConstrainPeriodicBoundaries(ConstraintBase): diff --git a/elastica/boundary_conditions.py b/elastica/boundary_conditions.py index 7f87258ea..3a1b99642 100644 --- a/elastica/boundary_conditions.py +++ b/elastica/boundary_conditions.py @@ -184,7 +184,7 @@ def constrain_rates(self, system: SystemType, time: np.floating) -> None: ) @staticmethod - @njit(cache=True) + @njit(cache=True) # type: ignore def compute_constrain_values( position_collection: NDArray[np.floating], fixed_position_collection: NDArray[np.floating], @@ -213,7 +213,7 @@ def compute_constrain_values( director_collection[..., 0] = fixed_directors_collection @staticmethod - @njit(cache=True) + @njit(cache=True) # type: ignore def compute_constrain_rates( velocity_collection: NDArray[np.floating], omega_collection: NDArray[np.floating], @@ -365,7 +365,7 @@ def constrain_rates(self, system: SystemType, time: np.floating) -> None: ) @staticmethod - @njit(cache=True) + @njit(cache=True) # type: ignore def nb_constrain_translational_values( position_collection: NDArray[np.floating], fixed_position_collection: NDArray[np.floating], @@ -403,7 +403,7 @@ def nb_constrain_translational_values( ] @staticmethod - @njit(cache=True) + @njit(cache=True) # type: ignore def nb_constrain_translational_rates( velocity_collection: NDArray[np.floating], indices: NDArray[np.integer], @@ -434,7 +434,7 @@ def nb_constrain_translational_rates( ) * velocity_collection[..., k] @staticmethod - @njit(cache=True) + @njit(cache=True) # type: ignore def nb_constrain_rotational_rates( director_collection: NDArray[np.floating], omega_collection: NDArray[np.floating], @@ -552,7 +552,7 @@ def constrain_rates(self, system: SystemType, time: np.floating) -> None: ) @staticmethod - @njit(cache=True) + @njit(cache=True) # type: ignore def nb_constraint_rotational_values( director_collection: NDArray[np.floating], fixed_director_collection: NDArray[np.floating], @@ -575,7 +575,7 @@ def nb_constraint_rotational_values( director_collection[..., k] = fixed_director_collection[..., i] @staticmethod - @njit(cache=True) + @njit(cache=True) # type: ignore def nb_constrain_translational_values( position_collection: NDArray[np.floating], fixed_position_collection: NDArray[np.floating], @@ -598,7 +598,7 @@ def nb_constrain_translational_values( position_collection[..., k] = fixed_position_collection[..., i] @staticmethod - @njit(cache=True) + @njit(cache=True) # type: ignore def nb_constrain_translational_rates( velocity_collection: NDArray[np.floating], indices: NDArray[np.integer] ) -> None: @@ -620,7 +620,7 @@ def nb_constrain_translational_rates( velocity_collection[2, k] = 0.0 @staticmethod - @njit(cache=True) + @njit(cache=True) # type: ignore def nb_constrain_rotational_rates( omega_collection: NDArray[np.floating], indices: NDArray[np.integer] ) -> None: diff --git a/elastica/dissipation.py b/elastica/dissipation.py index f0a79f560..dfe5373af 100644 --- a/elastica/dissipation.py +++ b/elastica/dissipation.py @@ -248,7 +248,7 @@ def dampen_rates(self, rod: RodType, time: np.floating) -> None: ) -@njit(cache=True) +@njit(cache=True) # type: ignore def _filter_function_periodic_condition_ring_rod( velocity_collection: NDArray[np.floating], velocity_filter_term: NDArray[np.floating], @@ -286,7 +286,7 @@ def _filter_function_periodic_condition_ring_rod( omega_collection[:] = omega_collection_with_periodic_bc[:, 1:-1] -@njit(cache=True) +@njit(cache=True) # type: ignore def _filter_function_periodic_condition( velocity_collection: NDArray[np.floating], velocity_filter_term: NDArray[np.floating], @@ -306,7 +306,7 @@ def _filter_function_periodic_condition( ) -@njit(cache=True) +@njit(cache=True) # type: ignore def nb_filter_rate( rate_collection: NDArray[np.floating], filter_term: NDArray[np.floating], diff --git a/elastica/experimental/connection_contact_joint/parallel_connection.py b/elastica/experimental/connection_contact_joint/parallel_connection.py index fe67a048d..757987482 100644 --- a/elastica/experimental/connection_contact_joint/parallel_connection.py +++ b/elastica/experimental/connection_contact_joint/parallel_connection.py @@ -125,7 +125,7 @@ def apply_forces(self, rod_one, index_one, rod_two, index_two): ) @staticmethod - @njit(cache=True) + @njit(cache=True) # type: ignore def _apply_forces( k, nu, @@ -262,7 +262,7 @@ def apply_torques(self, rod_one, index_one, rod_two, index_two): ) @staticmethod - @njit(cache=True) + @njit(cache=True) # type: ignore def _apply_torques( spring_force, rod_one_rd2, diff --git a/elastica/experimental/interaction.py b/elastica/experimental/interaction.py index 253688971..776f0da7e 100644 --- a/elastica/experimental/interaction.py +++ b/elastica/experimental/interaction.py @@ -66,7 +66,7 @@ def apply_forces(self, system, time=0.0): ) -@njit(cache=True) +@njit(cache=True) # type: ignore def anisotropic_friction_numba_rigid_body( plane_origin, plane_normal, diff --git a/elastica/external_forces.py b/elastica/external_forces.py index 62c4e6adb..00cbe60d0 100644 --- a/elastica/external_forces.py +++ b/elastica/external_forces.py @@ -92,7 +92,7 @@ def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: ) @staticmethod - @njit(cache=True) + @njit(cache=True) # type: ignore def compute_gravity_forces( acc_gravity: NDArray[np.floating], mass: NDArray[np.floating], @@ -166,7 +166,7 @@ def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: ) @staticmethod - @njit(cache=True) + @njit(cache=True) # type: ignore def compute_end_point_forces( external_forces: NDArray[np.floating], start_force: NDArray[np.floating], @@ -375,7 +375,7 @@ def apply_torques(self, rod: SystemType, time: np.floating = 0.0) -> None: ) @staticmethod - @njit(cache=True) + @njit(cache=True) # type: ignore def compute_muscle_torques( time: np.floating, my_spline: NDArray[np.floating], @@ -412,7 +412,7 @@ def compute_muscle_torques( ) -@njit(cache=True) +@njit(cache=True) # type: ignore def inplace_addition( external_force_or_torque: NDArray[np.floating], force_or_torque: NDArray[np.floating], @@ -438,7 +438,7 @@ def inplace_addition( external_force_or_torque[i, k] += force_or_torque[i, k] -@njit(cache=True) +@njit(cache=True) # type: ignore def inplace_substraction( external_force_or_torque: NDArray[np.floating], force_or_torque: NDArray[np.floating], diff --git a/elastica/interaction.py b/elastica/interaction.py index d695e6153..f0b9a9b28 100644 --- a/elastica/interaction.py +++ b/elastica/interaction.py @@ -45,7 +45,7 @@ def nodes_to_elements(input: Any) -> NoReturn: ) -@njit(cache=True) +@njit(cache=True) # type: ignore def elements_to_nodes_inplace( vector_in_element_frame: Any, vector_in_node_frame: Any ) -> NoReturn: @@ -316,7 +316,7 @@ def anisotropic_friction( # Slender body module -@njit(cache=True) +@njit(cache=True) # type: ignore def sum_over_elements(input: NDArray[np.floating]) -> np.floating: """ This function sums all elements of the input array. @@ -382,7 +382,7 @@ def node_to_element_pos_or_vel(vector_in_node_frame: Any) -> NoReturn: ) -@njit(cache=True) +@njit(cache=True) # type: ignore def slender_body_forces( tangents: NDArray[np.floating], velocity_collection: NDArray[np.floating], @@ -579,7 +579,7 @@ def apply_normal_force( ) -@njit(cache=True) +@njit(cache=True) # type: ignore def apply_normal_force_numba_rigid_body( plane_origin: Any, plane_normal: Any, diff --git a/elastica/reset_functions_for_block_structure/_reset_ghost_vector_or_scalar.py b/elastica/reset_functions_for_block_structure/_reset_ghost_vector_or_scalar.py index 7ed2a3bf4..3bac67448 100644 --- a/elastica/reset_functions_for_block_structure/_reset_ghost_vector_or_scalar.py +++ b/elastica/reset_functions_for_block_structure/_reset_ghost_vector_or_scalar.py @@ -3,7 +3,7 @@ from numba import njit -@njit(cache=True) +@njit(cache=True) # type: ignore def _reset_vector_ghost(input, ghost_idx, reset_value=0.0): """ This function resets the ghost of an input vector collection. Default reset value is 0.0. @@ -39,7 +39,7 @@ def _reset_vector_ghost(input, ghost_idx, reset_value=0.0): input[i, k] = reset_value -@njit(cache=True) +@njit(cache=True) # type: ignore def _reset_scalar_ghost(input, ghost_idx, reset_value=0.0): """ This function resets the ghost of a scalar collection. Default reset value is 0.0. diff --git a/elastica/rod/data_structures.py b/elastica/rod/data_structures.py index c65a9c4b5..c4accfd97 100644 --- a/elastica/rod/data_structures.py +++ b/elastica/rod/data_structures.py @@ -432,7 +432,7 @@ def __init__( self.director_collection = director_collection_view -@njit(cache=True) +@njit(cache=True) # type: ignore def overload_operator_kinematic_numba( n_nodes: int, prefac: np.floating, @@ -535,7 +535,7 @@ def dynamic_rates( return prefac * self.dvdt_dwdt_collection -@njit(cache=True) +@njit(cache=True) # type: ignore def overload_operator_dynamic_numba( rate_collection: NDArray[np.floating], scaled_second_deriv_array: NDArray[np.floating], diff --git a/elastica/rod/knot_theory.py b/elastica/rod/knot_theory.py index ef469257d..4f959ce06 100644 --- a/elastica/rod/knot_theory.py +++ b/elastica/rod/knot_theory.py @@ -191,7 +191,7 @@ def compute_twist( return total_twist, local_twist -@njit(cache=True) +@njit(cache=True) # type: ignore def _compute_twist( center_line: NDArray[np.floating], normal_collection: NDArray[np.floating] ) -> tuple[NDArray[np.floating], NDArray[np.floating]]: @@ -322,7 +322,7 @@ def compute_writhe( return total_writhe -@njit(cache=True) +@njit(cache=True) # type: ignore def _compute_writhe(center_line: NDArray[np.floating]) -> NDArray[np.floating]: """ Parameters @@ -478,7 +478,7 @@ def compute_link( return total_link -@njit(cache=True) +@njit(cache=True) # type: ignore def _compute_auxiliary_line( center_line: NDArray[np.floating], normal_collection: NDArray[np.floating], @@ -537,7 +537,7 @@ def _compute_auxiliary_line( return auxiliary_line -@njit(cache=True) +@njit(cache=True) # type: ignore def _compute_link( center_line: NDArray[np.floating], auxiliary_line: NDArray[np.floating] ) -> NDArray[np.floating]: @@ -617,7 +617,7 @@ def _compute_link( return total_link -@njit(cache=True) +@njit(cache=True) # type: ignore def _compute_auxiliary_line_added_segments( beginning_direction: NDArray[np.floating], end_direction: NDArray[np.floating], @@ -663,7 +663,7 @@ def _compute_auxiliary_line_added_segments( return new_auxiliary_line -@njit(cache=True) +@njit(cache=True) # type: ignore def _compute_additional_segment( center_line: NDArray[np.floating], segment_length: np.floating, From d7ba02e39fe2c11a35bdc5fac816a971c97d4292 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Thu, 13 Jun 2024 13:33:52 +0900 Subject: [PATCH 057/134] update: mypy pass for connection module --- elastica/modules/callbacks.py | 27 ++++++++----------- elastica/modules/connections.py | 48 ++++++++++++++++----------------- elastica/modules/protocol.py | 5 ++-- elastica/systems/protocol.py | 5 ++++ elastica/typing.py | 5 ++-- 5 files changed, 45 insertions(+), 45 deletions(-) diff --git a/elastica/modules/callbacks.py b/elastica/modules/callbacks.py index 472b042c4..9812b4ef2 100644 --- a/elastica/modules/callbacks.py +++ b/elastica/modules/callbacks.py @@ -4,10 +4,9 @@ Provides the callBack interface to collect data over time (see `callback_functions.py`). """ -from typing import Type, Protocol +from typing import Type, Protocol, Any from typing_extensions import Self # 3.11: from typing import Self from elastica.typing import SystemType, SystemIdxType, OperatorFinalizeType -from elastica.typing import CallbackParam from .protocol import ModuleProtocol import numpy as np @@ -25,11 +24,7 @@ class SystemCollectionWithCallBacksProtocol(SystemCollectionProtocol, Protocol): def collect_diagnostics(self, system: SystemType) -> ModuleProtocol: ... def _callback_execution( - self, - time: np.floating, - current_step: int, - *args: CallbackParam.args, - **kwargs: CallbackParam.kwargs, + self, time: np.floating, current_step: int, *args: Any, **kwargs: Any ) -> None: ... @@ -122,14 +117,14 @@ def __init__(self, sys_idx: SystemIdxType): """ self._sys_idx: SystemIdxType = sys_idx self._callback_cls: Type[CallBackBaseClass] - self._args: CallbackParam.args - self._kwargs: CallbackParam.kwargs + self._args: Any + self._kwargs: Any def using( self, - callback_cls: Type[CallBackBaseClass], - *args: CallbackParam.args, - **kwargs: CallbackParam.kwargs, + cls: Type[CallBackBaseClass], + *args: Any, + **kwargs: Any, ) -> Self: """ This method is a module to set which callback class is used to collect data @@ -137,7 +132,7 @@ def using( Parameters ---------- - callback_cls: object + cls: object User defined callback class. Returns @@ -145,11 +140,11 @@ def using( """ assert issubclass( - callback_cls, CallBackBaseClass + cls, CallBackBaseClass ), "{} is not a valid call back. Did you forget to derive from CallBackClass?".format( - callback_cls + cls ) - self._callback_cls = callback_cls + self._callback_cls = cls self._args = args self._kwargs = kwargs return self diff --git a/elastica/modules/connections.py b/elastica/modules/connections.py index 92966b782..11ede0d83 100644 --- a/elastica/modules/connections.py +++ b/elastica/modules/connections.py @@ -5,21 +5,21 @@ Provides the connections interface to connect entities (rods, rigid bodies) using joints (see `joints.py`). """ -from typing import Type, Protocol, TypeAlias, cast +from typing import Type, Protocol, TypeAlias, cast, Any from typing_extensions import Self from elastica.typing import ( SystemIdxType, - OperatorParam, OperatorFinalizeType, SystemType, ) import numpy as np + from elastica.joint import FreeJoint from .protocol import SystemCollectionProtocol, ModuleProtocol ConnectionIndex: TypeAlias = ( - int | np.int_ | list[int] | tuple[int] | np.ndarray[int] | None + int | np.int_ | list[int] | tuple[int] | np.typing.NDArray | None ) @@ -32,8 +32,8 @@ def connect( self, first_rod: SystemType, second_rod: SystemType, - first_connect_idx: int = None, - second_connect_idx: int = None, + first_connect_idx: ConnectionIndex, + second_connect_idx: ConnectionIndex, ) -> ModuleProtocol: ... @@ -50,7 +50,7 @@ class Connections: """ def __init__(self: SystemCollectionWithConnectionProtocol): - self._connections = [] + self._connections: list[ModuleProtocol] = [] super(Connections, self).__init__() self._feature_group_finalize.append(self._finalize_connections) @@ -81,17 +81,17 @@ def connect( ------- """ - sys_idx = [None] * 2 - for i_sys, sys in enumerate((first_rod, second_rod)): - sys_idx[i_sys] = self._get_sys_idx_if_valid(sys) + sys_idx_first = self._get_sys_idx_if_valid(first_rod) + sys_idx_second = self._get_sys_idx_if_valid(second_rod) # For each system identified, get max dofs - # FIXME: Revert back to len, it should be able to take, systems without elements! - # sys_dofs = [len(self._systems[idx]) for idx in sys_idx] - sys_dofs = [self._systems[idx].n_elems for idx in sys_idx] + sys_dofs_first = self._systems[sys_idx_first].n_elems + sys_dofs_second = self._systems[sys_idx_second].n_elems # Create _Connect object, cache it and return to user - _connect: ModuleProtocol = _Connect(*sys_idx, *sys_dofs) + _connect: ModuleProtocol = _Connect( + sys_idx_first, sys_idx_second, sys_dofs_first, sys_dofs_second + ) _connect.set_index(first_connect_idx, second_connect_idx) # type: ignore[attr-defined] self._connections.append(_connect) self._feature_group_synchronize.append_id(_connect) @@ -178,9 +178,9 @@ def __init__( self._second_sys_idx: SystemIdxType = second_sys_idx self._first_sys_n_lim: int = first_sys_nlim self._second_sys_n_lim: int = second_sys_nlim - self._connect_cls: Type[FreeJoint] = None - self.first_sys_connection_idx: ConnectionIndex = None - self.second_sys_connection_idx: ConnectionIndex = None + self._connect_cls: Type[FreeJoint] + self.first_sys_connection_idx: ConnectionIndex + self.second_sys_connection_idx: ConnectionIndex def set_index( self, first_idx: ConnectionIndex, second_idx: ConnectionIndex @@ -256,9 +256,9 @@ def set_index( def using( self, - connect_cls: Type[FreeJoint], - *args: OperatorParam.args, - **kwargs: OperatorParam.kwargs, + cls: Type[FreeJoint], + *args: Any, + **kwargs: Any, ) -> Self: """ This method is a module to set which joint class is used to connect @@ -266,8 +266,8 @@ def using( Parameters ---------- - connect_cls: object - User defined callback class. + cls: object + User defined connection class. *args Variable length argument list **kwargs @@ -278,11 +278,11 @@ def using( """ assert issubclass( - connect_cls, FreeJoint + cls, FreeJoint ), "{} is not a valid joint class. Did you forget to derive from FreeJoint?".format( - connect_cls + cls ) - self._connect_cls = connect_cls + self._connect_cls = cls self._args = args self._kwargs = kwargs return self diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index 9de78e0a2..930de5a00 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -1,4 +1,4 @@ -from typing import Protocol, Generator, TypeVar, Any +from typing import Protocol, Generator, TypeVar, Any, Type from typing_extensions import Self # 3.11: from typing import Self from elastica.typing import ( SystemIdxType, @@ -7,6 +7,7 @@ OperatorFinalizeType, SystemType, ) +from elastica.joint import FreeJoint from .operator_group import OperatorGroupFIFO @@ -38,7 +39,7 @@ def _get_sys_idx_if_valid(self, sys_to_be_added: SystemType) -> SystemIdxType: . class ModuleProtocol(Protocol[M]): - def using(self, callback_cls, *args, **kwargs) -> Self: ... + def using(self, cls: Type[M], *args: Any, **kwargs: Any) -> Self: ... def instantiate(self) -> M: ... diff --git a/elastica/systems/protocol.py b/elastica/systems/protocol.py index db88665ca..8c49d8920 100644 --- a/elastica/systems/protocol.py +++ b/elastica/systems/protocol.py @@ -20,6 +20,9 @@ class SystemProtocol(Protocol): @property def n_nodes(self) -> int: ... + @property + def n_elems(self) -> int: ... + @property def position_collection(self) -> NDArray: ... @@ -78,3 +81,5 @@ def __call__(self, time: np.floating, dt: np.floating) -> np.floating: ... def state(self) -> StateType: ... @state.setter def state(self, state: StateType) -> None: ... + @property + def n_elems(self) -> int: ... diff --git a/elastica/typing.py b/elastica/typing.py index 51e7b851a..c150864a0 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -108,7 +108,6 @@ # Operators in elastica.modules OperatorParam = ParamSpec("OperatorParam") -CallbackParam = ParamSpec("CallbackParam") OperatorType: TypeAlias = Callable[OperatorParam, None] -OperatorCallbackType: TypeAlias = Callable[CallbackParam, None] -OperatorFinalizeType: TypeAlias = Callable +OperatorCallbackType: TypeAlias = Callable[..., None] +OperatorFinalizeType: TypeAlias = Callable[..., None] From 5dcfe75d8f13c4e328c8e69e75f305b6ced75ad4 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Thu, 13 Jun 2024 14:02:26 +0900 Subject: [PATCH 058/134] remove unused protocol --- elastica/modules/callbacks.py | 23 +++++------------------ elastica/modules/connections.py | 22 ++++------------------ elastica/modules/protocol.py | 33 +++++++++++++++++++++++++++++++++ 3 files changed, 42 insertions(+), 36 deletions(-) diff --git a/elastica/modules/callbacks.py b/elastica/modules/callbacks.py index 9812b4ef2..f2c353d3f 100644 --- a/elastica/modules/callbacks.py +++ b/elastica/modules/callbacks.py @@ -4,7 +4,7 @@ Provides the callBack interface to collect data over time (see `callback_functions.py`). """ -from typing import Type, Protocol, Any +from typing import Type, Any from typing_extensions import Self # 3.11: from typing import Self from elastica.typing import SystemType, SystemIdxType, OperatorFinalizeType from .protocol import ModuleProtocol @@ -15,19 +15,6 @@ from .protocol import SystemCollectionProtocol -class SystemCollectionWithCallBacksProtocol(SystemCollectionProtocol, Protocol): - _callback_list: list[ModuleProtocol] - _callback_operators: list[tuple[int, CallBackBaseClass]] - - _finalize_callback: OperatorFinalizeType - - def collect_diagnostics(self, system: SystemType) -> ModuleProtocol: ... - - def _callback_execution( - self, time: np.floating, current_step: int, *args: Any, **kwargs: Any - ) -> None: ... - - class CallBacks: """ CallBacks class is a module for calling callback functions, set by the user. If the user @@ -40,7 +27,7 @@ class CallBacks: List of call back classes defined for rod-like objects. """ - def __init__(self: SystemCollectionWithCallBacksProtocol) -> None: + def __init__(self: SystemCollectionProtocol) -> None: self._callback_list: list[ModuleProtocol] = [] self._callback_operators: list[tuple[int, CallBackBaseClass]] = [] super(CallBacks, self).__init__() @@ -48,7 +35,7 @@ def __init__(self: SystemCollectionWithCallBacksProtocol) -> None: self._feature_group_finalize.append(self._finalize_callback) def collect_diagnostics( - self: SystemCollectionWithCallBacksProtocol, system: SystemType + self: SystemCollectionProtocol, system: SystemType ) -> ModuleProtocol: """ This method calls user-defined call-back classes for a @@ -72,7 +59,7 @@ def collect_diagnostics( return _callbacks - def _finalize_callback(self: SystemCollectionWithCallBacksProtocol) -> None: + def _finalize_callback(self: SystemCollectionProtocol) -> None: # dev : the first index stores the rod index to collect data. self._callback_operators = [ (callback.id(), callback.instantiate()) for callback in self._callback_list @@ -85,7 +72,7 @@ def _finalize_callback(self: SystemCollectionWithCallBacksProtocol) -> None: self._callback_execution(time=time, current_step=0) def _callback_execution( - self: SystemCollectionWithCallBacksProtocol, + self: SystemCollectionProtocol, time: np.floating, current_step: int, ) -> None: diff --git a/elastica/modules/connections.py b/elastica/modules/connections.py index 11ede0d83..f32ee96f0 100644 --- a/elastica/modules/connections.py +++ b/elastica/modules/connections.py @@ -5,7 +5,7 @@ Provides the connections interface to connect entities (rods, rigid bodies) using joints (see `joints.py`). """ -from typing import Type, Protocol, TypeAlias, cast, Any +from typing import Type, TypeAlias, cast, Any from typing_extensions import Self from elastica.typing import ( SystemIdxType, @@ -23,20 +23,6 @@ ) -class SystemCollectionWithConnectionProtocol(SystemCollectionProtocol, Protocol): - _connections: list[ModuleProtocol] - - _finalize_connections: OperatorFinalizeType - - def connect( - self, - first_rod: SystemType, - second_rod: SystemType, - first_connect_idx: ConnectionIndex, - second_connect_idx: ConnectionIndex, - ) -> ModuleProtocol: ... - - class Connections: """ The Connections class is a module for connecting rod-like objects using joints selected @@ -49,13 +35,13 @@ class Connections: List of joint classes defined for rod-like objects. """ - def __init__(self: SystemCollectionWithConnectionProtocol): + def __init__(self: SystemCollectionProtocol): self._connections: list[ModuleProtocol] = [] super(Connections, self).__init__() self._feature_group_finalize.append(self._finalize_connections) def connect( - self: SystemCollectionWithConnectionProtocol, + self: SystemCollectionProtocol, first_rod: SystemType, second_rod: SystemType, first_connect_idx: ConnectionIndex = None, @@ -98,7 +84,7 @@ def connect( return _connect - def _finalize_connections(self: SystemCollectionWithConnectionProtocol) -> None: + def _finalize_connections(self: SystemCollectionProtocol) -> None: # From stored _Connect objects, instantiate the joints and store it # dev : the first indices stores the # (first rod index, second_rod_idx, connection_idx_on_first_rod, connection_idx_on_second_rod) diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index 930de5a00..08d5fee05 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -1,5 +1,6 @@ from typing import Protocol, Generator, TypeVar, Any, Type from typing_extensions import Self # 3.11: from typing import Self +from abc import abstractmethod from elastica.typing import ( SystemIdxType, OperatorType, @@ -8,8 +9,10 @@ SystemType, ) from elastica.joint import FreeJoint +from elastica.callback_functions import CallBackBaseClass from .operator_group import OperatorGroupFIFO +from .connection import ConnectionIndex class SystemCollectionProtocol(Protocol): @@ -34,6 +37,36 @@ def blocks(self) -> Generator[SystemType, None, None]: ... def _get_sys_idx_if_valid(self, sys_to_be_added: SystemType) -> SystemIdxType: ... + # Connection API + _finalize_connections: OperatorFinalizeType + _connections: list[ModuleProtocol] + + @abstractmethod + def connect( + self, + first_rod: SystemType, + second_rod: SystemType, + first_connect_idx: ConnectionIndex, + second_connect_idx: ConnectionIndex, + ) -> ModuleProtocol: + raise NotImplementedError + + # CallBack API + _callback_list: list[ModuleProtocol] + _callback_operators: list[tuple[int, CallBackBaseClass]] + + _finalize_callback: OperatorFinalizeType + + @abstractmethod + def collect_diagnostics(self, system: SystemType) -> ModuleProtocol: + raise NotImplementedError + + @abstractmethod + def _callback_execution( + self, time: np.floating, current_step: int, *args: Any, **kwargs: Any + ) -> None: + raise NotImplementedError + M = TypeVar("M", bound="ModuleProtocol") From 2a08279948b5ce5a9c703f62f41bf40ff2297599 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Thu, 13 Jun 2024 14:18:09 +0900 Subject: [PATCH 059/134] typing: edit FreeJoint base indexing type --- elastica/joint.py | 43 ++++++++++++++++++--------------- elastica/modules/connections.py | 4 +-- 2 files changed, 25 insertions(+), 22 deletions(-) diff --git a/elastica/joint.py b/elastica/joint.py index 1e10eb5a8..2000bae6d 100644 --- a/elastica/joint.py +++ b/elastica/joint.py @@ -4,6 +4,7 @@ from elastica._rotations import _inv_rotate from elastica.typing import SystemType, RodType +from elastica.modules.connections import ConnectionIndex import numpy as np import logging from numpy.typing import NDArray @@ -48,9 +49,9 @@ def __init__(self, k: np.floating, nu: np.floating) -> None: def apply_forces( self, system_one: SystemType, - index_one: int, + index_one: ConnectionIndex, system_two: SystemType, - index_two: int, + index_two: ConnectionIndex, ) -> None: """ Apply joint force to the connected rod objects. @@ -59,11 +60,11 @@ def apply_forces( ---------- system_one : SystemType Rod or rigid-body object - index_one : int + index_one : int | np.ndarray Index of first rod for joint. system_two : SystemType Rod or rigid-body object - index_two : int + index_two : int | np.ndarray Index of second rod for joint. Returns @@ -91,9 +92,9 @@ def apply_forces( def apply_torques( self, system_one: SystemType, - index_one: int, + index_one: ConnectionIndex, system_two: SystemType, - index_two: int, + index_two: ConnectionIndex, ) -> None: """ Apply restoring joint torques to the connected rod objects. @@ -104,11 +105,11 @@ def apply_torques( ---------- system_one : SystemType Rod or rigid-body object - index_one : int + index_one : int | np.ndarray Index of first rod for joint. system_two : SystemType Rod or rigid-body object - index_two : int + index_two : int | np.ndarray Index of second rod for joint. Returns @@ -172,18 +173,18 @@ def __init__( def apply_forces( self, system_one: SystemType, - index_one: int, + index_one: ConnectionIndex, system_two: SystemType, - index_two: int, + index_two: ConnectionIndex, ) -> None: return super().apply_forces(system_one, index_one, system_two, index_two) def apply_torques( self, system_one: SystemType, - index_one: int, + index_one: ConnectionIndex, system_two: SystemType, - index_two: int, + index_two: ConnectionIndex, ) -> None: # current tangent direction of the `index_two` element of system two system_two_tangent = system_two.director_collection[2, :, index_two] @@ -279,18 +280,18 @@ def __init__( def apply_forces( self, system_one: SystemType, - index_one: int, + index_one: ConnectionIndex, system_two: SystemType, - index_two: int, + index_two: ConnectionIndex, ) -> None: return super().apply_forces(system_one, index_one, system_two, index_two) def apply_torques( self, system_one: SystemType, - index_one: int, + index_one: ConnectionIndex, system_two: SystemType, - index_two: int, + index_two: ConnectionIndex, ) -> None: # collect directors of systems one and two # note that systems can be either rods or rigid bodies @@ -618,9 +619,9 @@ def __init__( def apply_forces( self, rod_one: SystemType, - index_one: int, + index_one: ConnectionIndex, rod_two: SystemType, - index_two: int, + index_two: ConnectionIndex, ) -> None: # del index_one, index_two from elastica.contact_utils import ( @@ -737,7 +738,11 @@ def __init__(self, k: np.floating, nu: np.floating) -> None: ) def apply_forces( - self, rod_one: SystemType, index_one: int, rod_two: SystemType, index_two: int + self, + rod_one: SystemType, + index_one: ConnectionIndex, + rod_two: SystemType, + index_two: ConnectionIndex, ) -> None: # del index_one, index_two from elastica._contact_functions import ( diff --git a/elastica/modules/connections.py b/elastica/modules/connections.py index d77e1532b..b659ae324 100644 --- a/elastica/modules/connections.py +++ b/elastica/modules/connections.py @@ -97,7 +97,7 @@ def apply_forces_and_torques( first_connect_idx: ConnectionIndex, system_two: SystemType, second_connect_idx: ConnectionIndex, - ): + ) -> None: connect_instance.apply_forces( system_one=system_one, index_one=first_connect_idx, @@ -129,8 +129,6 @@ def apply_forces_and_torques( self._feature_group_synchronize.add_operators(connection, [func]) - self.warnings(connection) - self._connections = [] del self._connections From 4ab8e4c7e52c45fdf432f10c9b91d9ef69dcb6ed Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Thu, 13 Jun 2024 16:11:56 +0900 Subject: [PATCH 060/134] typing: finish constraints hinting --- elastica/modules/connections.py | 3 +- elastica/modules/constraints.py | 117 ++++++++++++++++++-------------- elastica/modules/protocol.py | 23 ++++++- elastica/systems/protocol.py | 3 + elastica/typing.py | 2 +- 5 files changed, 91 insertions(+), 57 deletions(-) diff --git a/elastica/modules/connections.py b/elastica/modules/connections.py index b659ae324..de7db3acc 100644 --- a/elastica/modules/connections.py +++ b/elastica/modules/connections.py @@ -18,6 +18,7 @@ from .protocol import SystemCollectionProtocol, ModuleProtocol +# TODO: Maybe just use slice?? ConnectionIndex: TypeAlias = ( int | np.int_ | list[int] | tuple[int] | np.typing.NDArray | None ) @@ -35,7 +36,7 @@ class Connections: List of joint classes defined for rod-like objects. """ - def __init__(self: SystemCollectionProtocol): + def __init__(self: SystemCollectionProtocol) -> None: self._connections: list[ModuleProtocol] = [] super(Connections, self).__init__() self._feature_group_finalize.append(self._finalize_connections) diff --git a/elastica/modules/constraints.py b/elastica/modules/constraints.py index 42c46fdd5..466ce118b 100644 --- a/elastica/modules/constraints.py +++ b/elastica/modules/constraints.py @@ -4,9 +4,19 @@ Provides the constraints interface to enforce displacement boundary conditions (see `boundary_conditions.py`). """ +from typing import Any, Type, TypeAlias +from typing_extensions import Self + +import numpy as np from elastica.boundary_conditions import ConstraintBase +from elastica.typing import SystemType, SystemIdxType +from .protocol import SystemCollectionProtocol, ModuleProtocol + +# TODO: Maybe just use slice?? +ConstrainingIndex: TypeAlias = list[int] | tuple[int] | np.typing.NDArray | None + class Constraints: """ @@ -20,14 +30,14 @@ class Constraints: List of boundary condition classes defined for rod-like objects. """ - def __init__(self): - self._constraints = [] + def __init__(self: SystemCollectionProtocol) -> None: + self._constraints_list: list[ModuleProtocol] = [] super(Constraints, self).__init__() self._feature_group_constrain_values.append(self._constrain_values) self._feature_group_constrain_rates.append(self._constrain_rates) self._feature_group_finalize.append(self._finalize_constraints) - def constrain(self, system): + def constrain(self: SystemCollectionProtocol, system: SystemType) -> ModuleProtocol: """ This method enforces a displacement boundary conditions to the relevant user-defined system or rod-like object. You must input the system or rod-like @@ -45,12 +55,12 @@ def constrain(self, system): sys_idx = self._get_sys_idx_if_valid(system) # Create _Constraint object, cache it and return to user - _constraint = _Constraint(sys_idx) - self._constraints.append(_constraint) + _constraint: ModuleProtocol = _Constraint(sys_idx) + self._constraints_list.append(_constraint) return _constraint - def _finalize_constraints(self): + def _finalize_constraints(self: SystemCollectionProtocol) -> None: """ In case memory block have ring rod, then periodic boundaries have to be synched. In order to synchronize periodic boundaries, a new constrain for memory block rod added called as _ConstrainPeriodicBoundaries. This @@ -74,9 +84,9 @@ def _finalize_constraints(self): # dev : the first index stores the rod index to apply the boundary condition # to. Technically we can use another array but it its one more book-keeping # step. Being lazy, I put them both in the same array - self._constraints[:] = [ - (constraint.id(), constraint(self._systems[constraint.id()])) - for constraint in self._constraints + self._constraints_operators = [ + (constraint.id(), constraint.instantiate(self._systems[constraint.id()])) + for constraint in self._constraints_list ] # Sort from lowest id to highest id for potentially better memory access @@ -85,20 +95,20 @@ def _finalize_constraints(self): # [(0, ConstraintBase, OneEndFixedBC), (1, HelicalBucklingBC), ... ] # Thus using lambda we iterate over the list of tuples and use rod number (x[0]) # to sort constraints. - self._constraints.sort(key=lambda x: x[0]) + self._constraints_operators.sort(key=lambda x: x[0]) # At t=0.0, constrain all the boundary conditions (for compatability with # initial conditions) self._constrain_values(time=0.0) self._constrain_rates(time=0.0) - def _constrain_values(self, time, *args, **kwargs): - for sys_id, constraint in self._constraints: - constraint.constrain_values(self._systems[sys_id], time, *args, **kwargs) + def _constrain_values(self: SystemCollectionProtocol, time: np.floating) -> None: + for sys_id, constraint in self._constraints_operators: + constraint.constrain_values(self._systems[sys_id], time) - def _constrain_rates(self, time, *args, **kwargs): - for sys_id, constraint in self._constraints: - constraint.constrain_rates(self._systems[sys_id], time, *args, **kwargs) + def _constrain_rates(self: SystemCollectionProtocol, time: np.floating) -> None: + for sys_id, constraint in self._constraints_operators: + constraint.constrain_rates(self._systems[sys_id], time) class _Constraint: @@ -108,14 +118,16 @@ class _Constraint: Attributes ---------- _sys_idx: int - _bc_cls: list + _bc_cls: Type[ConstraintBase] + constrained_position_idx: ConstrainingIndex + constrained_director_idx: ConstrainingIndex *args Variable length argument list. **kwargs Arbitrary keyword arguments. """ - def __init__(self, sys_idx: int): + def __init__(self, sys_idx: SystemIdxType) -> None: """ Parameters @@ -124,18 +136,27 @@ def __init__(self, sys_idx: int): """ self._sys_idx = sys_idx - self._bc_cls = None - self._args = () - self._kwargs = {} - - def using(self, bc_cls, *args, **kwargs): + self._bc_cls: Type[ConstraintBase] + self._args: Any + self._kwargs: Any + self.constrained_position_idx: ConstrainingIndex + self.constrained_director_idx: ConstrainingIndex + + def using( + self, + cls: Type[ConstraintBase], + constrained_position_idx: ConstrainingIndex = None, + constrained_director_idx: ConstrainingIndex = None, + *args: Any, + **kwargs: Any, + ) -> Self: """ This method is a module to set which boundary condition class is used to enforce boundary condition from user defined rod-like objects. Parameters ---------- - bc_cls : object + cls : Type[ConstraintBase] User defined boundary condition class. *args Variable length argument list @@ -147,61 +168,53 @@ def using(self, bc_cls, *args, **kwargs): """ assert issubclass( - bc_cls, ConstraintBase + cls, ConstraintBase ), "{} is not a valid constraint. Constraint must be driven from ConstraintBase.".format( - bc_cls + cls ) - self._bc_cls = bc_cls + self._bc_cls = cls + self.constrained_position_idx = constrained_position_idx + self.constrained_director_idx = constrained_director_idx self._args = args self._kwargs = kwargs return self - def id(self): + def id(self) -> SystemIdxType: return self._sys_idx - def __call__(self, rod, *args, **kwargs): - """Constructs a constraint after checks - - Parameters - ---------- - args - kwargs - - Returns - ------- - - """ + def instantiate(self, system: SystemType) -> ConstraintBase: + """Constructs a constraint after checks""" if not self._bc_cls: raise RuntimeError( "No boundary condition provided to constrain rod" "id {0} at {1}, but a BC was intended. Did you" - "forget to call the `using` method?".format(self.id(), rod) + "forget to call the `using` method?".format(self.id(), system) ) - # If there is position, director in kwargs, deal with it first - # Returns None if not found - pos_indices = self._kwargs.get( - "constrained_position_idx", None - ) # calculate position indices as a tuple director_indices = self._kwargs.get( "constrained_director_idx", None ) # calculate director indices as a tuple - # If pos_indices is not None, construct list else empty list # IMPORTANT : do copy for memory-safe operations positions = ( - [rod.position_collection[..., idx].copy() for idx in pos_indices] - if pos_indices + [ + system.position_collection[..., idx].copy() + for idx in self.constrained_position_idx + ] + if self.constrained_position_idx else [] ) directors = ( - [rod.director_collection[..., idx].copy() for idx in director_indices] - if director_indices + [ + system.director_collection[..., idx].copy() + for idx in self.constrained_director_idx + ] + if self.constrained_director_idx else [] ) try: bc = self._bc_cls( - *positions, *directors, *self._args, _system=rod, **self._kwargs + *positions, *directors, *self._args, _system=system, **self._kwargs ) return bc except (TypeError, IndexError): diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index 08d5fee05..3822601fe 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -10,6 +10,7 @@ ) from elastica.joint import FreeJoint from elastica.callback_functions import CallBackBaseClass +from elastica.boundary_conditions import ConstraintBase from .operator_group import OperatorGroupFIFO from .connection import ConnectionIndex @@ -52,11 +53,10 @@ def connect( raise NotImplementedError # CallBack API + _finalize_callback: OperatorFinalizeType _callback_list: list[ModuleProtocol] _callback_operators: list[tuple[int, CallBackBaseClass]] - _finalize_callback: OperatorFinalizeType - @abstractmethod def collect_diagnostics(self, system: SystemType) -> ModuleProtocol: raise NotImplementedError @@ -67,6 +67,23 @@ def _callback_execution( ) -> None: raise NotImplementedError + # Constraints API + _constraints_list: list[ModuleProtocol] + _constraints_operators: list[tuple[int, ConstraintBase]] + _finalize_constraints: OperatorFinalizeType + + @abstractmethod + def constrain(self, system: SystemType) -> ModuleProtocol: + raise NotImplementedError + + @abstractmethod + def _constrain_values(self, time: np.floating) -> None: + raise NotImplementedError + + @abstractmethod + def _constrain_rates(self, time: np.floating) -> None: + raise NotImplementedError + M = TypeVar("M", bound="ModuleProtocol") @@ -74,6 +91,6 @@ def _callback_execution( class ModuleProtocol(Protocol[M]): def using(self, cls: Type[M], *args: Any, **kwargs: Any) -> Self: ... - def instantiate(self) -> M: ... + def instantiate(self, system: SystemType) -> M: ... def id(self) -> Any: ... diff --git a/elastica/systems/protocol.py b/elastica/systems/protocol.py index 8c49d8920..6b23d3293 100644 --- a/elastica/systems/protocol.py +++ b/elastica/systems/protocol.py @@ -29,6 +29,9 @@ def position_collection(self) -> NDArray: ... @property def velocity_collection(self) -> NDArray: ... + @property + def director_collection(self) -> NDArray: ... + @property def acceleration_collection(self) -> NDArray: ... diff --git a/elastica/typing.py b/elastica/typing.py index c150864a0..7732b1409 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -56,7 +56,7 @@ FreeJoint = None -SystemType: TypeAlias = Union[SymplecticSystemProtocol, ExplicitSystemProtocol] +SystemType: TypeAlias = SymplecticSystemProtocol # | ExplicitSystemProtocol SystemIdxType: TypeAlias = int # ModuleObjectTypes: TypeAlias = ( From 9d75a896ebaf1fb4481bbd88dd6774de34df8df8 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Thu, 13 Jun 2024 16:51:23 +0900 Subject: [PATCH 061/134] typing: finish contact hinting --- elastica/modules/base_system.py | 5 +- elastica/modules/callbacks.py | 2 +- elastica/modules/connections.py | 2 +- elastica/modules/constraints.py | 7 +-- elastica/modules/contact.py | 93 ++++++++++++++++++--------------- elastica/modules/protocol.py | 17 +++++- elastica/typing.py | 34 +----------- 7 files changed, 76 insertions(+), 84 deletions(-) diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index 1bf549d1d..4e331443a 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -13,6 +13,7 @@ OperatorType, OperatorCallbackType, OperatorFinalizeType, + AllowedContactType, ) import numpy as np @@ -134,7 +135,9 @@ def override_allowed_types( self.allowed_sys_types = allowed_types @final - def _get_sys_idx_if_valid(self, sys_to_be_added: SystemType) -> SystemIdxType: + def _get_sys_idx_if_valid( + self, sys_to_be_added: SystemType | AllowedContactType + ) -> SystemIdxType: n_systems = len(self) # Total number of systems from mixed-in class sys_idx: SystemIdxType diff --git a/elastica/modules/callbacks.py b/elastica/modules/callbacks.py index f2c353d3f..f36d9cec2 100644 --- a/elastica/modules/callbacks.py +++ b/elastica/modules/callbacks.py @@ -139,7 +139,7 @@ def using( def id(self) -> SystemIdxType: return self._sys_idx - def instantiate(self) -> CallBackBaseClass: + def instantiate(self, system: SystemType) -> CallBackBaseClass: """Constructs a callback functions after checks""" if not self._callback_cls: raise RuntimeError( diff --git a/elastica/modules/connections.py b/elastica/modules/connections.py index de7db3acc..8344d0e0f 100644 --- a/elastica/modules/connections.py +++ b/elastica/modules/connections.py @@ -296,7 +296,7 @@ def id( self.second_sys_connection_idx, ) - def instantiate(self) -> FreeJoint: + def instantiate(self, system: SystemType) -> FreeJoint: if self._connect_cls is None: raise RuntimeError( "No connections provided to link rod id {0}" diff --git a/elastica/modules/constraints.py b/elastica/modules/constraints.py index 466ce118b..388b21119 100644 --- a/elastica/modules/constraints.py +++ b/elastica/modules/constraints.py @@ -82,8 +82,7 @@ def _finalize_constraints(self: SystemCollectionProtocol) -> None: # inplace : https://stackoverflow.com/a/1208792 # dev : the first index stores the rod index to apply the boundary condition - # to. Technically we can use another array but it its one more book-keeping - # step. Being lazy, I put them both in the same array + # to. self._constraints_operators = [ (constraint.id(), constraint.instantiate(self._systems[constraint.id()])) for constraint in self._constraints_list @@ -191,10 +190,6 @@ def instantiate(self, system: SystemType) -> ConstraintBase: "forget to call the `using` method?".format(self.id(), system) ) - director_indices = self._kwargs.get( - "constrained_director_idx", None - ) # calculate director indices as a tuple - # IMPORTANT : do copy for memory-safe operations positions = ( [ diff --git a/elastica/modules/contact.py b/elastica/modules/contact.py index 8882e92f1..ae24f7c2d 100644 --- a/elastica/modules/contact.py +++ b/elastica/modules/contact.py @@ -5,13 +5,30 @@ Provides the contact interface to apply contact forces between objects (rods, rigid bodies, surfaces). """ +from typing import Type, Any +from typing_extensions import Self +from elastica.typing import ( + SystemIdxType, + OperatorFinalizeType, + SystemType, + AllowedContactType, +) +from .protocol import SystemCollectionProtocol, ModuleProtocol + import logging import functools -from elastica.typing import SystemType, AllowedContactType + +import numpy as np + +from elastica.contact_forces import NoContact logger = logging.getLogger(__name__) +def warnings() -> None: + logger.warning("Contact features should be instantiated lastly.") + + class Contact: """ The Contact class is a module for applying contact between rod-like objects . To apply contact between rod-like objects, @@ -23,14 +40,16 @@ class Contact: List of contact classes defined for rod-like objects. """ - def __init__(self): - self._contacts = [] + def __init__(self: SystemCollectionProtocol) -> None: + self._contacts: list[ModuleProtocol] = [] super(Contact, self).__init__() self._feature_group_finalize.append(self._finalize_contact) def detect_contact_between( - self, first_system: SystemType, second_system: AllowedContactType - ): + self: SystemCollectionProtocol, + first_system: SystemType, + second_system: AllowedContactType, + ) -> ModuleProtocol: """ This method adds contact detection between two objects using the selected contact class. You need to input the two objects that are to be connected. @@ -46,31 +65,31 @@ def detect_contact_between( ------- """ - sys_idx = [None] * 2 - for i_sys, sys in enumerate((first_system, second_system)): - sys_idx[i_sys] = self._get_sys_idx_if_valid(sys) + sys_idx_first = self._get_sys_idx_if_valid(first_system) + sys_idx_second = self._get_sys_idx_if_valid(second_system) # Create _Contact object, cache it and return to user - _contact = _Contact(*sys_idx) + _contact = _Contact(sys_idx_first, sys_idx_second) self._contacts.append(_contact) self._feature_group_synchronize.append_id(_contact) return _contact - def _finalize_contact(self) -> None: + def _finalize_contact(self: SystemCollectionProtocol) -> None: # dev : the first indices stores the # (first_rod_idx, second_rod_idx) # to apply the contacts to - # Technically we can use another array but it its one more book-keeping - # step. Being lazy, I put them both in the same array def apply_contact( - time, contact_instance, system, first_sys_idx, second_sys_idx - ): + time: np.floating, + contact_instance: NoContact, + first_sys_idx: SystemIdxType, + second_sys_idx: SystemIdxType, + ) -> None: contact_instance.apply_contact( - system_one=system[first_sys_idx], - system_two=system[second_sys_idx], + system_one=self._systems[first_sys_idx], + system_two=self._systems[second_sys_idx], ) for contact in self._contacts: @@ -84,25 +103,17 @@ def apply_contact( func = functools.partial( apply_contact, contact_instance=contact_instance, - system=self._systems, first_sys_idx=first_sys_idx, second_sys_idx=second_sys_idx, ) self._feature_group_synchronize.add_operators(contact, [func]) - self.warnings(contact) + if not self._feature_group_synchronize.is_last(contact): + warnings() self._contacts = [] del self._contacts - def warnings(self, contact): - from elastica.contact_forces import NoContact - - # Classes that should be used last - if not self._feature_group_synchronize.is_last(contact): - if isinstance(contact._contact_cls, NoContact): - logger.warning("Contact features should be instantiated lastly.") - class _Contact: """ @@ -110,9 +121,9 @@ class _Contact: Attributes ---------- - _first_sys_idx: int - _second_sys_idx: int - _contact_cls: list + _first_sys_idx: SystemIdxType + _second_sys_idx: SystemIdxType + _contact_cls: Type[NoContact] *args Variable length argument list. **kwargs @@ -121,8 +132,8 @@ class _Contact: def __init__( self, - first_sys_idx: int, - second_sys_idx: int, + first_sys_idx: SystemIdxType, + second_sys_idx: SystemIdxType, ) -> None: """ @@ -133,16 +144,18 @@ def __init__( """ self.first_sys_idx = first_sys_idx self.second_sys_idx = second_sys_idx - self._contact_cls = None + self._contact_cls: Type[NoContact] + self._args: Any + self._kwargs: Any - def using(self, contact_cls: object, *args, **kwargs): + def using(self, cls: Type[NoContact], *args: Any, **kwargs: Any) -> Self: """ This method is a module to set which contact class is used to apply contact between user defined rod-like objects. Parameters ---------- - contact_cls: object + cls: Type[NoContact] User defined contact class. *args Variable length argument list @@ -153,25 +166,23 @@ def using(self, contact_cls: object, *args, **kwargs): ------- """ - from elastica.contact_forces import NoContact - assert issubclass( - contact_cls, NoContact + cls, NoContact ), "{} is not a valid contact class. Did you forget to derive from NoContact?".format( - contact_cls + cls ) - self._contact_cls = contact_cls + self._contact_cls = cls self._args = args self._kwargs = kwargs return self - def id(self): + def id(self) -> Any: return ( self.first_sys_idx, self.second_sys_idx, ) - def instantiate(self, *args, **kwargs): + def instantiate(self) -> NoContact: if not self._contact_cls: raise RuntimeError( "No contacts provided to to establish contact between rod-like object id {0}" diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index 3822601fe..a2ffd44ec 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -7,6 +7,7 @@ OperatorCallbackType, OperatorFinalizeType, SystemType, + AllowedContactType, ) from elastica.joint import FreeJoint from elastica.callback_functions import CallBackBaseClass @@ -36,7 +37,9 @@ def _feature_group_finalize(self) -> list[OperatorFinalizeType]: ... def blocks(self) -> Generator[SystemType, None, None]: ... - def _get_sys_idx_if_valid(self, sys_to_be_added: SystemType) -> SystemIdxType: ... + def _get_sys_idx_if_valid( + self, sys_to_be_added: SystemType | AllowedContactType + ) -> SystemIdxType: ... # Connection API _finalize_connections: OperatorFinalizeType @@ -84,6 +87,16 @@ def _constrain_values(self, time: np.floating) -> None: def _constrain_rates(self, time: np.floating) -> None: raise NotImplementedError + # Contact API + _contacts: list[ModuleProtocol] + _finalize_contact: OperatorFinalizeType + + @abstractmethod + def detect_contact_between( + self, first_system: SystemType, second_system: AllowedContactType + ) -> ModuleProtocol: + raise NotImplementedError + M = TypeVar("M", bound="ModuleProtocol") @@ -91,6 +104,6 @@ def _constrain_rates(self, time: np.floating) -> None: class ModuleProtocol(Protocol[M]): def using(self, cls: Type[M], *args: Any, **kwargs: Any) -> Self: ... - def instantiate(self, system: SystemType) -> M: ... + def instantiate(self, *args: Any, **kwargs: Any) -> M: ... def id(self) -> Any: ... diff --git a/elastica/typing.py b/elastica/typing.py index 7732b1409..5b357b4cc 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -4,7 +4,7 @@ """ from typing import TYPE_CHECKING -from typing import Type, Union, Callable, Any, ParamSpec +from typing import Type, Callable, Any, ParamSpec from typing import TypeAlias @@ -67,44 +67,14 @@ # Mostly used in explicit stepper: for symplectic, use kinetic and dynamic state StateType: TypeAlias = State -# NoOpt stepper -# Symplectic stepper -# StepOperatorType = Callable[ -# [SymplecticStepperProtocol, SymplecticSystemProtocol, np.floating, np.floating], None -# ] -# PrefactorOperatorType = Callable[ -# [SymplecticStepperProtocol, np.floating], np.floating -# ] OperatorType: TypeAlias = Callable[ ..., Any ] # TODO: Maybe can be more specific. Up for discussion. SteppersOperatorsType: TypeAlias = tuple[tuple[OperatorType, ...], ...] -# tuple[Union[PrefactorOperatorType, StepOperatorType, NoOpType, np.floating], ...], ... -# Explicit stepper -# ExplicitStageOperatorType = Callable[ -# [ -# SymplecticStepperProtocol, -# ExplicitSystemProtocol, -# MemoryProtocol, -# np.floating, -# np.floating, -# ], -# None, -# ] -# ExplicitUpdateOperatorType = Callable[ -# [ -# SymplecticStepperProtocol, -# ExplicitSystemProtocol, -# MemoryProtocol, -# np.floating, -# np.floating, -# ], -# np.floating, -# ] RodType: TypeAlias = Type[RodBase] SystemCollectionType: TypeAlias = SystemCollectionProtocol -AllowedContactType: TypeAlias = Union[SystemType, Type[SurfaceBase]] +AllowedContactType: TypeAlias = SystemType | Type[SurfaceBase] # Operators in elastica.modules OperatorParam = ParamSpec("OperatorParam") From 842451f7c98bac3be3f6316317c26d185ee385ee Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Thu, 13 Jun 2024 17:58:07 +0900 Subject: [PATCH 062/134] test: update unittest for refactor codes --- elastica/joint.py | 3 +- elastica/modules/base_system.py | 8 +-- elastica/modules/callbacks.py | 4 +- elastica/modules/connections.py | 20 +++----- elastica/modules/constraints.py | 11 ++--- elastica/modules/contact.py | 2 +- elastica/modules/protocol.py | 34 ++++++++----- elastica/timestepper/__init__.py | 2 +- elastica/typing.py | 67 +++++++++++++++----------- tests/test_modules/test_base_system.py | 4 +- tests/test_modules/test_constraints.py | 20 ++++---- 11 files changed, 94 insertions(+), 81 deletions(-) diff --git a/elastica/joint.py b/elastica/joint.py index 2000bae6d..260df2724 100644 --- a/elastica/joint.py +++ b/elastica/joint.py @@ -3,8 +3,7 @@ from typing import Any, NoReturn, Optional from elastica._rotations import _inv_rotate -from elastica.typing import SystemType, RodType -from elastica.modules.connections import ConnectionIndex +from elastica.typing import SystemType, RodType, ConnectionIndex import numpy as np import logging from numpy.typing import NDArray diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index 4e331443a..8a0d17c78 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -104,17 +104,17 @@ def _check_type(self, sys_to_be_added: SystemType) -> bool: def __len__(self) -> int: return len(self._systems) - def __getitem__(self, idx: int | slice, /) -> SystemType | list[SystemType]: # type: ignore + def __getitem__(self, idx, /): # type: ignore return self._systems[idx] - def __delitem__(self, idx: int | slice, /) -> None: # type: ignore + def __delitem__(self, idx, /): # type: ignore del self._systems[idx] - def __setitem__(self, idx: int | slice, system: SystemType | Iterable[SystemType], /) -> None: # type: ignore + def __setitem__(self, idx, system, /): # type: ignore self._check_type(system) self._systems[idx] = system - def insert(self, idx: int, system: SystemType) -> None: # type: ignore + def insert(self, idx, system) -> None: # type: ignore self._check_type(system) self._systems.insert(idx, system) diff --git a/elastica/modules/callbacks.py b/elastica/modules/callbacks.py index f36d9cec2..bf67357f3 100644 --- a/elastica/modules/callbacks.py +++ b/elastica/modules/callbacks.py @@ -139,9 +139,9 @@ def using( def id(self) -> SystemIdxType: return self._sys_idx - def instantiate(self, system: SystemType) -> CallBackBaseClass: + def instantiate(self) -> CallBackBaseClass: """Constructs a callback functions after checks""" - if not self._callback_cls: + if not hasattr(self, "_callback_cls"): raise RuntimeError( "No callback provided to act on rod id {0}" "but a callback was registered. Did you forget to call" diff --git a/elastica/modules/connections.py b/elastica/modules/connections.py index 8344d0e0f..a9ffb19d8 100644 --- a/elastica/modules/connections.py +++ b/elastica/modules/connections.py @@ -5,12 +5,13 @@ Provides the connections interface to connect entities (rods, rigid bodies) using joints (see `joints.py`). """ -from typing import Type, TypeAlias, cast, Any +from typing import Type, cast, Any from typing_extensions import Self from elastica.typing import ( SystemIdxType, OperatorFinalizeType, SystemType, + ConnectionIndex, ) import numpy as np import functools @@ -18,11 +19,6 @@ from .protocol import SystemCollectionProtocol, ModuleProtocol -# TODO: Maybe just use slice?? -ConnectionIndex: TypeAlias = ( - int | np.int_ | list[int] | tuple[int] | np.typing.NDArray | None -) - class Connections: """ @@ -149,8 +145,8 @@ class _Connect: _first_sys_n_lim: int _second_sys_n_lim: int _connect_class: list - first_sys_connection_idx: int - second_sys_connection_idx: int + first_sys_connection_idx: ConnectionIndex + second_sys_connection_idx: ConnectionIndex *args Variable length argument list. **kwargs @@ -177,9 +173,9 @@ def __init__( self._second_sys_idx: SystemIdxType = second_sys_idx self._first_sys_n_lim: int = first_sys_nlim self._second_sys_n_lim: int = second_sys_nlim + self.first_sys_connection_idx: ConnectionIndex = None + self.second_sys_connection_idx: ConnectionIndex = None self._connect_cls: Type[FreeJoint] - self.first_sys_connection_idx: ConnectionIndex - self.second_sys_connection_idx: ConnectionIndex def set_index( self, first_idx: ConnectionIndex, second_idx: ConnectionIndex @@ -296,8 +292,8 @@ def id( self.second_sys_connection_idx, ) - def instantiate(self, system: SystemType) -> FreeJoint: - if self._connect_cls is None: + def instantiate(self) -> FreeJoint: + if not hasattr(self, "_connect_cls"): raise RuntimeError( "No connections provided to link rod id {0}" "(at {2}) and {1} (at {3}), but a Connection" diff --git a/elastica/modules/constraints.py b/elastica/modules/constraints.py index 388b21119..1d95f3565 100644 --- a/elastica/modules/constraints.py +++ b/elastica/modules/constraints.py @@ -4,19 +4,16 @@ Provides the constraints interface to enforce displacement boundary conditions (see `boundary_conditions.py`). """ -from typing import Any, Type, TypeAlias +from typing import Any, Type from typing_extensions import Self import numpy as np from elastica.boundary_conditions import ConstraintBase -from elastica.typing import SystemType, SystemIdxType +from elastica.typing import SystemType, SystemIdxType, ConstrainingIndex from .protocol import SystemCollectionProtocol, ModuleProtocol -# TODO: Maybe just use slice?? -ConstrainingIndex: TypeAlias = list[int] | tuple[int] | np.typing.NDArray | None - class Constraints: """ @@ -144,9 +141,9 @@ def __init__(self, sys_idx: SystemIdxType) -> None: def using( self, cls: Type[ConstraintBase], + *args: Any, constrained_position_idx: ConstrainingIndex = None, constrained_director_idx: ConstrainingIndex = None, - *args: Any, **kwargs: Any, ) -> Self: """ @@ -183,7 +180,7 @@ def id(self) -> SystemIdxType: def instantiate(self, system: SystemType) -> ConstraintBase: """Constructs a constraint after checks""" - if not self._bc_cls: + if not hasattr(self, "_bc_cls"): raise RuntimeError( "No boundary condition provided to constrain rod" "id {0} at {1}, but a BC was intended. Did you" diff --git a/elastica/modules/contact.py b/elastica/modules/contact.py index ae24f7c2d..924e7a4a9 100644 --- a/elastica/modules/contact.py +++ b/elastica/modules/contact.py @@ -183,7 +183,7 @@ def id(self) -> Any: ) def instantiate(self) -> NoContact: - if not self._contact_cls: + if not hasattr(self, "_contact_cls"): raise RuntimeError( "No contacts provided to to establish contact between rod-like object id {0}" " and {1}, but a Contact" diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index a2ffd44ec..3e70c8bd8 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -8,13 +8,26 @@ OperatorFinalizeType, SystemType, AllowedContactType, + ConnectionIndex, ) from elastica.joint import FreeJoint from elastica.callback_functions import CallBackBaseClass from elastica.boundary_conditions import ConstraintBase +import numpy as np + from .operator_group import OperatorGroupFIFO -from .connection import ConnectionIndex + + +M = TypeVar("M", bound="ModuleProtocol") + + +class ModuleProtocol(Protocol[M]): + def using(self, cls: Type[M], *args: Any, **kwargs: Any) -> Self: ... + + def instantiate(self, *args: Any, **kwargs: Any) -> M: ... + + def id(self) -> Any: ... class SystemCollectionProtocol(Protocol): @@ -23,15 +36,23 @@ class SystemCollectionProtocol(Protocol): @property def _feature_group_synchronize(self) -> OperatorGroupFIFO: ... + def synchronize(self, time: np.floating) -> None: ... + @property def _feature_group_constrain_values(self) -> list[OperatorType]: ... + def constrain_values(self, time: np.floating) -> None: ... + @property def _feature_group_constrain_rates(self) -> list[OperatorType]: ... + def constrain_rates(self, time: np.floating) -> None: ... + @property def _feature_group_callback(self) -> list[OperatorCallbackType]: ... + def apply_callbacks(self, time: np.floating, current_step: int) -> None: ... + @property def _feature_group_finalize(self) -> list[OperatorFinalizeType]: ... @@ -96,14 +117,3 @@ def detect_contact_between( self, first_system: SystemType, second_system: AllowedContactType ) -> ModuleProtocol: raise NotImplementedError - - -M = TypeVar("M", bound="ModuleProtocol") - - -class ModuleProtocol(Protocol[M]): - def using(self, cls: Type[M], *args: Any, **kwargs: Any) -> Self: ... - - def instantiate(self, *args: Any, **kwargs: Any) -> M: ... - - def id(self) -> Any: ... diff --git a/elastica/timestepper/__init__.py b/elastica/timestepper/__init__.py index 34575dfd6..0fd32c751 100644 --- a/elastica/timestepper/__init__.py +++ b/elastica/timestepper/__init__.py @@ -55,7 +55,7 @@ def integrate( def integrate( stepper: StepperProtocol, - systems: SystemType | SystemCollectionType, + systems: "SystemType | SystemCollectionType", final_time: float, n_steps: int = 1000, restart_time: float = 0.0, diff --git a/elastica/typing.py b/elastica/typing.py index 5b357b4cc..4febe1065 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -7,6 +7,8 @@ from typing import Type, Callable, Any, ParamSpec from typing import TypeAlias +import numpy as np + if TYPE_CHECKING: # Used for type hinting without circular imports @@ -32,31 +34,31 @@ from .dissipation import DamperBase from .external_forces import NoForces from .joint import FreeJoint -else: - RodBase = None - RigidBodyBase = None - SurfaceBase = None - - SystemCollectionProtocol = None - - State = "State" - SymplecticSystemProtocol = None - ExplicitSystemProtocol = None - - StepperProtocol = None - SymplecticStepperProtocol = None - MemoryProtocol = None - - # Modules Base Classes - FreeBC = None - CallBackBaseClass = None - NoContact = None - DamperBase = None - NoForces = None - FreeJoint = None - - -SystemType: TypeAlias = SymplecticSystemProtocol # | ExplicitSystemProtocol +# else: +# RodBase = None +# RigidBodyBase = None +# SurfaceBase = None +# +# SystemCollectionProtocol = None +# +# State = "State" +# SymplecticSystemProtocol = None +# ExplicitSystemProtocol = None +# +# StepperProtocol = None +# SymplecticStepperProtocol = None +# MemoryProtocol = None +# +# # Modules Base Classes +# FreeBC = None +# CallBackBaseClass = None +# NoContact = None +# DamperBase = None +# NoForces = None +# FreeJoint = None + + +SystemType: TypeAlias = "SymplecticSystemProtocol" # | ExplicitSystemProtocol SystemIdxType: TypeAlias = int # ModuleObjectTypes: TypeAlias = ( @@ -65,16 +67,23 @@ # TODO: Modify this line and move to elastica/typing.py once system state is defined # Mostly used in explicit stepper: for symplectic, use kinetic and dynamic state -StateType: TypeAlias = State +StateType: TypeAlias = "State" OperatorType: TypeAlias = Callable[ ..., Any ] # TODO: Maybe can be more specific. Up for discussion. SteppersOperatorsType: TypeAlias = tuple[tuple[OperatorType, ...], ...] -RodType: TypeAlias = Type[RodBase] -SystemCollectionType: TypeAlias = SystemCollectionProtocol -AllowedContactType: TypeAlias = SystemType | Type[SurfaceBase] +RodType: TypeAlias = Type["RodBase"] +SystemCollectionType: TypeAlias = "SystemCollectionProtocol" +AllowedContactType: TypeAlias = SystemType | Type["SurfaceBase"] + +# Indexing types +# TODO: Maybe just use slice?? +ConstrainingIndex: TypeAlias = list[int] | tuple[int] | np.typing.NDArray | None +ConnectionIndex: TypeAlias = ( + int | np.int_ | list[int] | tuple[int] | np.typing.NDArray | None +) # Operators in elastica.modules OperatorParam = ParamSpec("OperatorParam") diff --git a/tests/test_modules/test_base_system.py b/tests/test_modules/test_base_system.py index e26791d66..e3a74dcd8 100644 --- a/tests/test_modules/test_base_system.py +++ b/tests/test_modules/test_base_system.py @@ -183,7 +183,9 @@ def test_constraint(self, load_collection, legal_constraint): simulator_class.constrain(rod).using(legal_constraint) simulator_class.finalize() # After finalize check if the created constrain object is instance of the class we have given. - assert isinstance(simulator_class._constraints[-1][-1], legal_constraint) + assert isinstance( + simulator_class._constraints_operators[-1][-1], legal_constraint + ) # TODO: this is a dummy test for constrain values and rates find a better way to test them simulator_class.constrain_values(time=0) diff --git a/tests/test_modules/test_constraints.py b/tests/test_modules/test_constraints.py index aeee1873b..069479d1a 100644 --- a/tests/test_modules/test_constraints.py +++ b/tests/test_modules/test_constraints.py @@ -45,7 +45,7 @@ def test_call_without_setting_constraint_throws_runtime_error( constraint = load_constraint with pytest.raises(RuntimeError) as excinfo: - constraint(None) # None is the rod/system parameter + constraint.instantiate(None) # None is the rod/system parameter assert "No boundary condition" in str(excinfo.value) def test_call_without_position_director_kwargs(self, load_constraint): @@ -60,7 +60,7 @@ def mock_init(self, *args, **kwargs): constraint.using(MockBC, 3.9, 4.0, "5", k=1, l_var="2", j=3.0) # Actual test is here, this should not throw - mock_bc = constraint(None) # None is Fake rod + mock_bc = constraint.instantiate(None) # None is Fake rod # More tests reinforcing the first assert mock_bc.dummy_one == 3.9 @@ -93,7 +93,7 @@ def mock_init(self, *args, **kwargs): # Actual test is here, this should not throw mock_rod = self.MockRod() - mock_bc = constraint(mock_rod) + mock_bc = constraint.instantiate(mock_rod) # More tests reinforcing the first for pos_idx_in_rod, pos_idx_in_bc in zip(position_indices, range(3)): @@ -125,7 +125,7 @@ def mock_init(self, *args, **kwargs): # Actual test is here, this should not throw mock_rod = self.MockRod() - mock_bc = constraint(mock_rod) + mock_bc = constraint.instantiate(mock_rod) # More tests reinforcing the first for dir_idx_in_rod, dir_idx_in_bc in zip(director_indices, range(3)): @@ -160,7 +160,7 @@ def mock_init(self, *args, **kwargs): # Actual test is here, this should not throw mock_rod = self.MockRod() - mock_bc = constraint(mock_rod) + mock_bc = constraint.instantiate(mock_rod) # More tests reinforcing the first pos_dir_offset = len(dof_indices) @@ -202,7 +202,7 @@ def mock_init(self, nu, **kwargs): mock_rod = self.MockRod() # Actual test is here, this should not throw with pytest.raises(TypeError) as excinfo: - _ = constraint(mock_rod) + _ = constraint.instantiate(mock_rod) assert "Unable to construct" in str(excinfo.value) @@ -281,7 +281,7 @@ def test_constrain_registers_and_returns_Constraint( scwc.append(mock_rod) _mock_constraint = scwc.constrain(mock_rod) - assert _mock_constraint in scwc._constraints + assert _mock_constraint in scwc._constraints_list assert _mock_constraint.__class__ == _Constraint from elastica.boundary_conditions import ConstraintBase @@ -318,7 +318,7 @@ def test_constrain_finalize_correctness(self, load_rod_with_constraints): scwc._finalize_constraints() - for x, y in scwc._constraints: + for x, y in scwc._constraints_operators: assert type(x) is int assert type(y) is bc_cls @@ -327,7 +327,7 @@ def test_constraint_properties(self, load_rod_with_constraints): scwc._finalize_constraints() for i in [0, 1, -1]: - x, y = scwc._constraints[i] + x, y = scwc._constraints_operators[i] mock_rod = scwc._systems[i] # Test system assert type(x) is int @@ -346,7 +346,7 @@ def test_constrain_finalize_sorted(self, load_rod_with_constraints): # this is allowed to fail (not critical) num = -np.inf - for x, _ in scwc._constraints: + for x, _ in scwc._constraints_list: assert num < x num = x From b4d54d2d40d044ed88dd57501a83a7c877cc8d8d Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Thu, 13 Jun 2024 18:14:19 +0900 Subject: [PATCH 063/134] typing: damping module finish --- elastica/modules/damping.py | 65 +++++++++++++++++------------- elastica/modules/protocol.py | 14 +++++++ tests/test_modules/test_damping.py | 12 +++--- 3 files changed, 56 insertions(+), 35 deletions(-) diff --git a/elastica/modules/damping.py b/elastica/modules/damping.py index 027eebdad..a5806d668 100644 --- a/elastica/modules/damping.py +++ b/elastica/modules/damping.py @@ -9,7 +9,14 @@ """ +from typing import Any, Type, List +from typing_extensions import Self + +import numpy as np + from elastica.dissipation import DamperBase +from elastica.typing import SystemType, SystemIdxType +from .protocol import SystemCollectionProtocol, ModuleProtocol class Damping: @@ -24,13 +31,13 @@ class Damping: List of damper classes defined for rod-like objects. """ - def __init__(self): - self._dampers = [] - super(Damping, self).__init__() + def __init__(self: SystemCollectionProtocol) -> None: + self._damping_list: List[ModuleProtocol] = [] + super().__init__() self._feature_group_constrain_rates.append(self._dampen_rates) self._feature_group_finalize.append(self._finalize_dampers) - def dampen(self, system): + def dampen(self: SystemCollectionProtocol, system: SystemType) -> ModuleProtocol: """ This method applies damping on relevant user-defined system or rod-like object. You must input the system or rod-like @@ -48,18 +55,18 @@ def dampen(self, system): sys_idx = self._get_sys_idx_if_valid(system) # Create _Damper object, cache it and return to user - _damper = _Damper(sys_idx) - self._dampers.append(_damper) + _damper: ModuleProtocol = _Damper(sys_idx) + self._damping_list.append(_damper) return _damper - def _finalize_dampers(self): + def _finalize_dampers(self: SystemCollectionProtocol) -> None: # From stored _Damping objects, instantiate the dissipation/damping # inplace : https://stackoverflow.com/a/1208792 - self._dampers[:] = [ - (damper.id(), damper(self._systems[damper.id()])) - for damper in self._dampers + self._damping_operators = [ + (damper.id(), damper.instantiate(self._systems[damper.id()])) + for damper in self._damping_list ] # Sort from lowest id to highest id for potentially better memory access @@ -67,11 +74,11 @@ def _finalize_dampers(self): # following elements are the type of damping. # Thus using lambda we iterate over the list of tuples and use rod number (x[0]) # to sort dampers. - self._dampers.sort(key=lambda x: x[0]) + self._damping_operators.sort(key=lambda x: x[0]) - def _dampen_rates(self, time, *args, **kwargs): - for sys_id, damper in self._dampers: - damper.dampen_rates(self._systems[sys_id], time, *args, **kwargs) + def _dampen_rates(self: SystemCollectionProtocol, time: np.floating) -> None: + for sys_id, damper in self._damping_operators: + damper.dampen_rates(self._systems[sys_id], time) class _Damper: @@ -88,7 +95,7 @@ class _Damper: Arbitrary keyword arguments. """ - def __init__(self, sys_idx: int): + def __init__(self, sys_idx: SystemIdxType) -> None: """ Parameters @@ -97,22 +104,22 @@ def __init__(self, sys_idx: int): """ self._sys_idx = sys_idx - self._damper_cls = None - self._args = () - self._kwargs = {} + self._damper_cls: Type[DamperBase] + self._args: Any + self._kwargs: Any - def using(self, damper_cls, *args, **kwargs): + def using(self, cls: Type[DamperBase], *args: Any, **kwargs: Any) -> Self: """ This method is a module to set which damper class is used to enforce damping from user defined rod-like objects. Parameters ---------- - damper_cls : object + cls : Type[DamperBase] User defined damper class. - *args - Variable length argument list - **kwargs + *args: Any + Variable length argument list. + **kwargs: Any Arbitrary keyword arguments. Returns @@ -120,19 +127,19 @@ def using(self, damper_cls, *args, **kwargs): """ assert issubclass( - damper_cls, DamperBase + cls, DamperBase ), "{} is not a valid damper. Damper must be driven from DamperBase.".format( - damper_cls + cls ) - self._damper_cls = damper_cls + self._damper_cls = cls self._args = args self._kwargs = kwargs return self - def id(self): + def id(self) -> SystemIdxType: return self._sys_idx - def __call__(self, rod, *args, **kwargs): + def instantiate(self, rod: SystemType) -> DamperBase: """Constructs a Damper class object after checks Parameters @@ -144,7 +151,7 @@ def __call__(self, rod, *args, **kwargs): ------- """ - if not self._damper_cls: + if not hasattr(self, "_damper_cls"): raise RuntimeError( "No damper provided to dampen rod id {0} at {1}," "but damping was intended. Did you" diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index 3e70c8bd8..2c4342895 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -13,6 +13,7 @@ from elastica.joint import FreeJoint from elastica.callback_functions import CallBackBaseClass from elastica.boundary_conditions import ConstraintBase +from elastica.dissipation import DamperBase import numpy as np @@ -117,3 +118,16 @@ def detect_contact_between( self, first_system: SystemType, second_system: AllowedContactType ) -> ModuleProtocol: raise NotImplementedError + + # Damping API + _damping_list: list[ModuleProtocol] + _damping_operators: list[tuple[int, DamperBase]] + _finalize_dampers: OperatorFinalizeType + + @abstractmethod + def dampen(self, system: SystemType) -> ModuleProtocol: + raise NotImplementedError + + @abstractmethod + def _dampen_rates(self, time: np.floating) -> None: + raise NotImplementedError diff --git a/tests/test_modules/test_damping.py b/tests/test_modules/test_damping.py index d36f0187c..2f18a5d64 100644 --- a/tests/test_modules/test_damping.py +++ b/tests/test_modules/test_damping.py @@ -39,7 +39,7 @@ def test_call_without_setting_damper_throws_runtime_error(self, load_damper): damper = load_damper with pytest.raises(RuntimeError) as excinfo: - damper(None) # None is the rod/system parameter + damper.instantiate(None) # None is the rod/system parameter assert "No damper" in str(excinfo.value) def test_call_with_args_and_kwargs(self, load_damper): @@ -56,7 +56,7 @@ def mock_init(self, *args, **kwargs): damper.using(MockDamper, 3.9, 4.0, "5", k=1, l_var="2", j=3.0) # Actual test is here, this should not throw - mock_damper = damper(None) # None is Fake rod + mock_damper = damper.instantiate(None) # None is Fake rod # More tests reinforcing the first assert mock_damper.dummy_one == 3.9 @@ -78,7 +78,7 @@ def test_call_improper_bc_throws_type_error(self, load_damper): mock_rod = self.MockRod() # Actual test is here, this should not throw with pytest.raises(TypeError) as excinfo: - _ = damper(mock_rod) + _ = damper.instantiate(mock_rod) assert "Unable to construct" in str(excinfo.value) @@ -153,7 +153,7 @@ def test_dampen_registers_and_returns_Damper(self, load_system_with_dampers): scwd.append(mock_rod) _mock_damper = scwd.dampen(mock_rod) - assert _mock_damper in scwd._dampers + assert _mock_damper in scwd._damping_list assert _mock_damper.__class__ == _Damper from elastica.dissipation import DamperBase @@ -185,7 +185,7 @@ def test_dampen_finalize_correctness(self, load_rod_with_dampers): scwd._finalize_dampers() - for x, y in scwd._dampers: + for x, y in scwd._damping_operators: assert type(x) is int assert type(y) is damper_cls @@ -194,7 +194,7 @@ def test_damper_properties(self, load_rod_with_dampers): scwd._finalize_dampers() for i in [0, 1, -1]: - x, y = scwd._dampers[i] + x, y = scwd._damping_operators[i] mock_rod = scwd._systems[i] # Test system assert type(x) is int From d32007d33c53420f0b527bdefbb501c44077cc04 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Thu, 13 Jun 2024 20:27:20 +0900 Subject: [PATCH 064/134] typing: finish forcing module hinting --- elastica/modules/damping.py | 12 +----- elastica/modules/forcing.py | 80 ++++++++++++++++-------------------- elastica/modules/protocol.py | 8 ++++ 3 files changed, 44 insertions(+), 56 deletions(-) diff --git a/elastica/modules/damping.py b/elastica/modules/damping.py index a5806d668..4fa0cf055 100644 --- a/elastica/modules/damping.py +++ b/elastica/modules/damping.py @@ -140,17 +140,7 @@ def id(self) -> SystemIdxType: return self._sys_idx def instantiate(self, rod: SystemType) -> DamperBase: - """Constructs a Damper class object after checks - - Parameters - ---------- - args - kwargs - - Returns - ------- - - """ + """Constructs a Damper class object after checks""" if not hasattr(self, "_damper_cls"): raise RuntimeError( "No damper provided to dampen rod id {0} at {1}," diff --git a/elastica/modules/forcing.py b/elastica/modules/forcing.py index 4d1c0d2d5..d6475d8a2 100644 --- a/elastica/modules/forcing.py +++ b/elastica/modules/forcing.py @@ -7,6 +7,14 @@ """ import logging import functools +from typing import Any, Type, List +from typing_extensions import Self + +import numpy as np + +from elastica.external_forces import NoForces +from elastica.typing import SystemType, SystemIdxType +from .protocol import SystemCollectionProtocol, ModuleProtocol logger = logging.getLogger(__name__) @@ -23,12 +31,14 @@ class Forcing: List of forcing class defined for rod-like objects. """ - def __init__(self): - self._ext_forces_torques = [] - super(Forcing, self).__init__() + def __init__(self: SystemCollectionProtocol) -> None: + self._ext_forces_torques: List[ModuleProtocol] = [] + super().__init__() self._feature_group_finalize.append(self._finalize_forcing) - def add_forcing_to(self, system): + def add_forcing_to( + self: SystemCollectionProtocol, system: SystemType + ) -> ModuleProtocol: """ This method applies external forces and torques on the relevant user-defined system or rod-like object. You must input the system @@ -52,7 +62,7 @@ def add_forcing_to(self, system): return _ext_force_torque - def _finalize_forcing(self): + def _finalize_forcing(self: SystemCollectionProtocol) -> None: # From stored _ExtForceTorque objects, and instantiate a Force # inplace : https://stackoverflow.com/a/1208792 @@ -73,14 +83,9 @@ def _finalize_forcing(self): external_force_and_torque, [apply_forces, apply_torques] ) - self.warnings(external_force_and_torque) - self._ext_forces_torques = [] del self._ext_forces_torques - def warnings(self, external_force_and_torque): - pass - class _ExtForceTorque: """ @@ -89,73 +94,58 @@ class _ExtForceTorque: Attributes ---------- _sys_idx: int - _forcing_cls: list - *args + _forcing_cls: Type[NoForces] + *args: Any Variable length argument list. - **kwargs + **kwargs: Any Arbitrary keyword arguments. """ - def __init__(self, sys_idx: int): + def __init__(self, sys_idx: SystemIdxType) -> None: """ - Parameters ---------- sys_idx: int """ self._sys_idx = sys_idx - self._forcing_cls = None - self._args = () - self._kwargs = {} + self._forcing_cls: Type[NoForces] + self._args: Any + self._kwargs: Any - def using(self, forcing_cls, *args, **kwargs): + def using(self, cls: Type[NoForces], *args: Any, **kwargs: Any) -> Self: """ - This method is a module to set which forcing class is used to apply forcing + This method sets which forcing class is used to apply forcing to user defined rod-like objects. Parameters ---------- - forcing_cls: object + cls: Type[Any] User defined forcing class. - *args - Variable length argument list - **kwargs + *args: Any + Variable length argument list. + **kwargs: Any Arbitrary keyword arguments. Returns ------- """ - from elastica.external_forces import NoForces - assert issubclass( - forcing_cls, NoForces + cls, NoForces ), "{} is not a valid forcing. Did you forget to derive from NoForces?".format( - forcing_cls + cls ) - self._forcing_cls = forcing_cls + self._forcing_cls = cls self._args = args self._kwargs = kwargs return self - def id(self): + def id(self) -> SystemIdxType: return self._sys_idx - def instantiate(self): - """Constructs a constraint after checks - - Parameters - ---------- - *args - Variable length argument list. - **kwargs - Arbitrary keyword arguments. - - Returns - ------- - - """ - if not self._forcing_cls: + def instantiate(self) -> NoForces: + """Constructs a constraint after checks""" + if not hasattr(self, "_forcing_cls"): raise RuntimeError( "No forcing provided to act on rod id {0}" "but a force was registered. Did you forget to call" diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index 2c4342895..5ce5ba7cf 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -109,6 +109,14 @@ def _constrain_values(self, time: np.floating) -> None: def _constrain_rates(self, time: np.floating) -> None: raise NotImplementedError + # Forcing API + _ext_forces_torques: list[ModuleProtocol] + _finalize_forcing: OperatorFinalizeType + + @abstractmethod + def add_forcing_to(self, system: SystemType) -> ModuleProtocol: + raise NotImplementedError + # Contact API _contacts: list[ModuleProtocol] _finalize_contact: OperatorFinalizeType From 1b0f109c72d319890cebfe21545308bba0f35e41 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Thu, 13 Jun 2024 21:34:11 +0900 Subject: [PATCH 065/134] typing: make memory block hinting agree to block module --- .../memory_block/memory_block_rigid_body.py | 24 ++++++---- elastica/memory_block/memory_block_rod.py | 36 +++++++-------- .../memory_block/memory_block_rod_base.py | 20 ++++----- elastica/modules/base_system.py | 13 ++---- elastica/modules/constraints.py | 4 +- elastica/modules/memory_block.py | 18 +++++--- elastica/rigidbody/rigid_body.py | 45 ++++++------------- elastica/rod/data_structures.py | 20 ++++----- elastica/rod/rod_base.py | 4 +- elastica/systems/protocol.py | 8 ++-- elastica/typing.py | 9 ++-- 11 files changed, 89 insertions(+), 112 deletions(-) diff --git a/elastica/memory_block/memory_block_rigid_body.py b/elastica/memory_block/memory_block_rigid_body.py index 4dd7fac4f..f6728c6f8 100644 --- a/elastica/memory_block/memory_block_rigid_body.py +++ b/elastica/memory_block/memory_block_rigid_body.py @@ -1,13 +1,17 @@ __doc__ = """Create block-structure class for collection of rigid body systems.""" import numpy as np -from typing import Sequence, Literal +from typing import Literal + +from elastica.typing import SystemType, SystemIdxType from elastica.rigidbody import RigidBodyBase from elastica.rigidbody.data_structures import _RigidRodSymplecticStepperMixin class MemoryBlockRigidBody(RigidBodyBase, _RigidRodSymplecticStepperMixin): - def __init__(self, systems: Sequence, system_idx_list: Sequence[np.int64]): + def __init__( + self, systems: list[RigidBodyBase], system_idx_list: list[SystemIdxType] + ): self.n_bodies = len(systems) self.n_elems = self.n_bodies @@ -23,7 +27,7 @@ def __init__(self, systems: Sequence, system_idx_list: Sequence[np.int64]): # Initialize the mixin class for symplectic time-stepper. _RigidRodSymplecticStepperMixin.__init__(self) - def _allocate_block_variables_scalars(self, systems: Sequence): + def _allocate_block_variables_scalars(self, systems: list[RigidBodyBase]): """ This function takes system collection and allocates the variables for block-structure and references allocated variables back to the systems. @@ -58,7 +62,7 @@ def _allocate_block_variables_scalars(self, systems: Sequence): value_type="scalar", ) - def _allocate_block_variables_vectors(self, systems: Sequence): + def _allocate_block_variables_vectors(self, systems: list[RigidBodyBase]): """ This function takes system collection and allocates the vector variables for block-structure and references allocated vector variables back to the systems. @@ -94,7 +98,7 @@ def _allocate_block_variables_vectors(self, systems: Sequence): value_type="vector", ) - def _allocate_block_variables_matrix(self, systems: Sequence): + def _allocate_block_variables_matrix(self, systems: list[RigidBodyBase]): """ This function takes system collection and allocates the matrix variables for block-structure and references allocated matrix variables back to the systems. @@ -130,7 +134,9 @@ def _allocate_block_variables_matrix(self, systems: Sequence): value_type="tensor", ) - def _allocate_block_variables_for_symplectic_stepper(self, systems: Sequence): + def _allocate_block_variables_for_symplectic_stepper( + self, systems: list[RigidBodyBase] + ): """ This function takes system collection and allocates the variables used by symplectic stepper for block-structure and references allocated variables back to the systems. @@ -176,7 +182,7 @@ def _allocate_block_variables_for_symplectic_stepper(self, systems: Sequence): def _map_system_properties_to_block_memory( self, mapping_dict: dict, - systems: Sequence, + systems: list[RigidBodyBase], block_memory: np.ndarray, value_type: Literal["scalar", "vector", "tensor"], ) -> None: @@ -186,8 +192,8 @@ def _map_system_properties_to_block_memory( ---------- mapping_dict: dict Dictionary with attribute names as keys and block row index as values. - systems: Sequence - A sequence containing Cosserat rod objects to map from. + systems: list[RigidBodyBase] + A sequence containing rigid body objects to map from. block_memory: ndarray Memory block that, at the end of the method execution, contains all designated attributes of all systems. diff --git a/elastica/memory_block/memory_block_rod.py b/elastica/memory_block/memory_block_rod.py index b6a6b73f7..fc7ba17bc 100644 --- a/elastica/memory_block/memory_block_rod.py +++ b/elastica/memory_block/memory_block_rod.py @@ -1,8 +1,8 @@ __doc__ = """Create block-structure class for collection of Cosserat rod systems.""" import numpy as np -from typing import Sequence, Literal, Callable +from typing import Literal, Callable +from elastica.typing import SystemIdxType from elastica.memory_block.memory_block_rod_base import ( - MemoryBlockRodBase, make_block_memory_metadata, make_block_memory_periodic_boundary_metadata, ) @@ -12,6 +12,8 @@ CosseratRod, _compute_sigma_kappa_for_blockstructure, ) +from elastica.rod.rod_base import RodBase +from elastica.rod.rod_base import RodBase from elastica._synchronize_periodic_boundary import ( _synchronize_periodic_boundary_of_vector_collection, _synchronize_periodic_boundary_of_scalar_collection, @@ -19,9 +21,7 @@ ) -class MemoryBlockCosseratRod( - MemoryBlockRodBase, CosseratRod, _RodSymplecticStepperMixin -): +class MemoryBlockCosseratRod(CosseratRod, _RodSymplecticStepperMixin): """ Memory block class for Cosserat rod equations. This class is derived from Cosserat Rod class in order to inherit the methods of Cosserat rod class. This class takes the cosserat rod object (systems) and creates big @@ -31,7 +31,9 @@ class MemoryBlockCosseratRod( TODO: need more documentation! """ - def __init__(self, systems: Sequence, system_idx_list): + def __init__( + self, systems: list[RodBase], system_idx_list: list[SystemIdxType] + ) -> None: # separate straight and ring rods system_straight_rod = [] @@ -200,7 +202,7 @@ def __init__(self, systems: Sequence, system_idx_list): # Initialize the mixin class for symplectic time-stepper. _RodSymplecticStepperMixin.__init__(self) - def _allocate_block_variables_in_nodes(self, systems: Sequence): + def _allocate_block_variables_in_nodes(self, systems: list[RodBase]) -> None: """ This function takes system collection and allocates the variables on node for block-structure and references allocated variables back to the @@ -250,7 +252,7 @@ def _allocate_block_variables_in_nodes(self, systems: Sequence): value_type="vector", ) - def _allocate_block_variables_in_elements(self, systems: Sequence): + def _allocate_block_variables_in_elements(self, systems: list[RodBase]) -> None: """ This function takes system collection and allocates the variables on elements for block-structure and references allocated variables back to the @@ -341,7 +343,7 @@ def _allocate_block_variables_in_elements(self, systems: Sequence): value_type="tensor", ) - def _allocate_blocks_variables_in_voronoi(self, systems: Sequence): + def _allocate_blocks_variables_in_voronoi(self, systems: list[RodBase]) -> None: """ This function takes system collection and allocates the variables on voronoi for block-structure and references allocated variables back to the @@ -408,18 +410,12 @@ def _allocate_blocks_variables_in_voronoi(self, systems: Sequence): value_type="tensor", ) - def _allocate_blocks_variables_for_symplectic_stepper(self, systems: Sequence): + def _allocate_blocks_variables_for_symplectic_stepper( + self, systems: list[RodBase] + ) -> None: """ This function takes system collection and allocates the variables used by symplectic stepper for block-structure and references allocated variables back to the systems. - - Parameters - ---------- - systems - - Returns - ------- - """ # These vectors are on nodes or on elements, but we stack them together for # better memory access. Because we use them together in time-steppers. @@ -475,7 +471,7 @@ def _allocate_blocks_variables_for_symplectic_stepper(self, systems: Sequence): def _map_system_properties_to_block_memory( self, mapping_dict: dict, - systems: Sequence, + systems: list[RodBase], block_memory: np.ndarray, domain_type: Literal["node", "element", "voronoi"], value_type: Literal["scalar", "vector", "tensor"], @@ -490,7 +486,7 @@ def _map_system_properties_to_block_memory( ---------- mapping_dict: dict Dictionary with attribute names as keys and block row index as values. - systems: Sequence + systems: list[RodBase] A sequence containing Cosserat rod objects to map from. block_memory: ndarray Memory block that, at the end of the method execution, contains all designated diff --git a/elastica/memory_block/memory_block_rod_base.py b/elastica/memory_block/memory_block_rod_base.py index c9bcc77e1..166a10c65 100644 --- a/elastica/memory_block/memory_block_rod_base.py +++ b/elastica/memory_block/memory_block_rod_base.py @@ -1,9 +1,12 @@ __doc__ = """Create block-structure class for collection of Cosserat rod systems.""" import numpy as np -from typing import Iterable +pass -def make_block_memory_metadata(n_elems_in_rods: np.ndarray) -> Iterable: + +def make_block_memory_metadata( + n_elems_in_rods: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """ This function, takes number of elements of each rod as a numpy array and computes, ghost nodes, elements and voronoi element indexes and numbers and returns it. @@ -52,7 +55,9 @@ def make_block_memory_metadata(n_elems_in_rods: np.ndarray) -> Iterable: return n_elems_with_ghosts, ghost_nodes_idx, ghost_elems_idx, ghost_voronoi_idx -def make_block_memory_periodic_boundary_metadata(n_elems_in_rods): +def make_block_memory_periodic_boundary_metadata( + n_elems_in_rods: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """ This function, takes the number of elements of ring rods and computes the periodic boundary node, element and voronoi index. @@ -153,12 +158,3 @@ def make_block_memory_periodic_boundary_metadata(n_elems_in_rods): periodic_boundary_elems_idx, periodic_boundary_voronoi_idx, ) - - -class MemoryBlockRodBase: - """ - This is the base class for memory blocks for rods. - """ - - def __init__(self): - pass diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index 8a0d17c78..4b816b1d6 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -5,7 +5,7 @@ Basic coordinating for multiple, smaller systems that have an independently integrable interface (i.e. works with symplectic or explicit routines `timestepper.py`.) """ -from typing import Type, Generator, Iterable +from typing import Type, Generator, Iterable, Any from typing import final from elastica.typing import ( SystemType, @@ -41,9 +41,6 @@ class BaseSystemCollection(MutableSequence): _systems: list List of rod-like objects. - Developer Note - ----- - Note ---- We can directly subclass a list for the @@ -64,7 +61,7 @@ def __init__(self) -> None: super().__init__() # List of system types/bases that are allowed - self.allowed_sys_types: tuple[Type[SystemType], ...] = ( + self.allowed_sys_types: tuple[Type, ...] = ( RodBase, RigidBodyBase, SurfaceBase, @@ -79,7 +76,7 @@ def __init__(self) -> None: self._finalize_flag: bool = False @final - def _check_type(self, sys_to_be_added: SystemType) -> bool: + def _check_type(self, sys_to_be_added: Any) -> bool: if not isinstance(sys_to_be_added, self.allowed_sys_types): raise TypeError( "{0}\n" @@ -135,9 +132,7 @@ def override_allowed_types( self.allowed_sys_types = allowed_types @final - def _get_sys_idx_if_valid( - self, sys_to_be_added: SystemType | AllowedContactType - ) -> SystemIdxType: + def _get_sys_idx_if_valid(self, sys_to_be_added: SystemType) -> SystemIdxType: n_systems = len(self) # Total number of systems from mixed-in class sys_idx: SystemIdxType diff --git a/elastica/modules/constraints.py b/elastica/modules/constraints.py index 1d95f3565..9c70e640e 100644 --- a/elastica/modules/constraints.py +++ b/elastica/modules/constraints.py @@ -95,8 +95,8 @@ def _finalize_constraints(self: SystemCollectionProtocol) -> None: # At t=0.0, constrain all the boundary conditions (for compatability with # initial conditions) - self._constrain_values(time=0.0) - self._constrain_rates(time=0.0) + self._constrain_values(time=np.float64(0.0)) + self._constrain_rates(time=np.float64(0.0)) def _constrain_values(self: SystemCollectionProtocol, time: np.floating) -> None: for sys_id, constraint in self._constraints_operators: diff --git a/elastica/modules/memory_block.py b/elastica/modules/memory_block.py index 6139e7eee..a12b70d71 100644 --- a/elastica/modules/memory_block.py +++ b/elastica/modules/memory_block.py @@ -3,13 +3,15 @@ Cosserat Rods, Rigid Body etc. """ +from elastica.typing import SystemType, SystemIdxType + from elastica.rod import RodBase from elastica.rigidbody import RigidBodyBase from elastica.surface import SurfaceBase from elastica.memory_block import MemoryBlockCosseratRod, MemoryBlockRigidBody -def construct_memory_block_structures(systems): +def construct_memory_block_structures(systems: list[SystemType]) -> list[SystemType]: """ This function takes the systems (rod or rigid body) appended to the simulator class and separates them into lists depending on if system is Cosserat rod or rigid body. Then using @@ -19,11 +21,11 @@ def construct_memory_block_structures(systems): ------- """ - _memory_blocks = [] - temp_list_for_cosserat_rod_systems = [] - temp_list_for_rigid_body_systems = [] - temp_list_for_cosserat_rod_systems_idx = [] - temp_list_for_rigid_body_systems_idx = [] + _memory_blocks: list[SystemType] = [] + temp_list_for_cosserat_rod_systems: list[RodBase] = [] + temp_list_for_rigid_body_systems: list[RigidBodyBase] = [] + temp_list_for_cosserat_rod_systems_idx: list[SystemIdxType] = [] + temp_list_for_rigid_body_systems_idx: list[SystemIdxType] = [] for system_idx, sys_to_be_added in enumerate(systems): @@ -36,7 +38,9 @@ def construct_memory_block_structures(systems): temp_list_for_rigid_body_systems_idx.append(system_idx) elif isinstance(sys_to_be_added, SurfaceBase): - pass + raise NotImplementedError( + "Surfaces are not yet implemented in memory block construction." + ) else: raise TypeError( diff --git a/elastica/rigidbody/rigid_body.py b/elastica/rigidbody/rigid_body.py index e87392744..7ef340d64 100644 --- a/elastica/rigidbody/rigid_body.py +++ b/elastica/rigidbody/rigid_body.py @@ -19,38 +19,19 @@ class RigidBodyBase(ABC): def __init__(self): - self.position_collection = NotImplementedError - self.velocity_collection = NotImplementedError - self.acceleration_collection = NotImplementedError - self.omega_collection = NotImplementedError - self.alpha_collection = NotImplementedError - self.director_collection = NotImplementedError - - self.external_forces = NotImplementedError - self.external_torques = NotImplementedError - - self.mass = NotImplementedError - - self.mass_second_moment_of_inertia = NotImplementedError - self.inv_mass_second_moment_of_inertia = NotImplementedError - - # @abstractmethod - # # def update_accelerations(self): - # # pass - - # def _compute_internal_forces_and_torques(self): - # """ - # This function here is only for integrator to work properly. We do not need - # internal forces and torques at all. - # Parameters - # ---------- - # time - # - # Returns - # ------- - # - # """ - # pass + self.position_collection: NDArray[np.floating] + self.velocity_collection: NDArray[np.floating] + self.omega_collection: NDArray[np.floating] + self.acceleration_collection: NDArray[np.floating] + self.director_collection: NDArray[np.floating] + self.alpha_collection: NDArray[np.floating] + self.external_forces: NDArray[np.floating] + self.external_torques: NDArray[np.floating] + + self.mass: np.floating + + self.mass_second_moment_of_inertia: NDArray[np.floating] + self.inv_mass_second_moment_of_inertia: NDArray[np.floating] def update_accelerations(self, time): np.copyto( diff --git a/elastica/rod/data_structures.py b/elastica/rod/data_structures.py index c4accfd97..35e2520a9 100644 --- a/elastica/rod/data_structures.py +++ b/elastica/rod/data_structures.py @@ -1,6 +1,6 @@ __doc__ = "Data structure wrapper for rod components" -from typing import Any, Optional +from typing import Optional from typing_extensions import Self import numpy as np from numpy.typing import NDArray @@ -63,20 +63,18 @@ def __init__(self) -> None: # is another function self.kinematic_rates = self.dynamic_states.kinematic_rates - def update_internal_forces_and_torques( - self, time: np.floating, *args: Any, **kwargs: Any - ) -> None: + def update_internal_forces_and_torques(self, time: np.floating) -> None: self.compute_internal_forces_and_torques(time) def dynamic_rates( - self, time: np.floating, prefac: np.floating, *args: Any, **kwargs: Any + self, + time: np.floating, + prefac: np.floating, ) -> NDArray[np.floating]: self.update_accelerations(time) - return self.dynamic_states.dynamic_rates(time, prefac, *args, **kwargs) + return self.dynamic_states.dynamic_rates(time, prefac) - def reset_external_forces_and_torques( - self, time: np.floating, *args: Any, **kwargs: Any - ) -> None: + def reset_external_forces_and_torques(self, time: np.floating) -> None: self.zeroed_out_external_forces_and_torques(time) @@ -519,9 +517,7 @@ def kinematic_rates( # Comes from kin_state -> (x,Q) += dt * (v,w) <- First part of dyn_state return self.velocity_collection, self.omega_collection - def dynamic_rates( - self, time: np.floating, prefac: np.floating - ) -> NDArray[np.floating]: + def dynamic_rates(self, time: np.floating, prefac: np.floating): """Yields dynamic rates to add to with _DynamicState Returns ------- diff --git a/elastica/rod/rod_base.py b/elastica/rod/rod_base.py index d49c0431f..dab11221f 100644 --- a/elastica/rod/rod_base.py +++ b/elastica/rod/rod_base.py @@ -22,8 +22,10 @@ def __init__(self) -> None: RodBase does not take any arguments. """ self.position_collection: NDArray[np.floating] - self.omega_collection: NDArray[np.floating] + self.velocity_collection: NDArray[np.floating] self.acceleration_collection: NDArray[np.floating] + self.director_collection: NDArray[np.floating] + self.omega_collection: NDArray[np.floating] self.alpha_collection: NDArray[np.floating] self.external_forces: NDArray[np.floating] self.external_torques: NDArray[np.floating] diff --git a/elastica/systems/protocol.py b/elastica/systems/protocol.py index 6b23d3293..5855e0bad 100644 --- a/elastica/systems/protocol.py +++ b/elastica/systems/protocol.py @@ -62,18 +62,18 @@ def kinematic_states(self) -> _KinematicState: ... def dynamic_states(self) -> _DynamicState: ... @property - def rate_collection(self) -> NDArray: ... + def rate_collection(self) -> NDArray[np.floating]: ... @property - def dvdt_dwdt_collection(self) -> NDArray: ... + def dvdt_dwdt_collection(self) -> NDArray[np.floating]: ... def kinematic_rates( self, time: np.floating, prefac: np.floating - ) -> tuple[NDArray, NDArray]: ... + ) -> tuple[NDArray[np.floating], NDArray[np.floating]]: ... def dynamic_rates( self, time: np.floating, prefac: np.floating - ) -> tuple[NDArray]: ... + ) -> NDArray[np.floating]: ... class ExplicitSystemProtocol(SystemProtocol, Protocol): diff --git a/elastica/typing.py b/elastica/typing.py index 4febe1065..b0f1f57af 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -69,14 +69,15 @@ # Mostly used in explicit stepper: for symplectic, use kinetic and dynamic state StateType: TypeAlias = "State" -OperatorType: TypeAlias = Callable[ - ..., Any -] # TODO: Maybe can be more specific. Up for discussion. +# TODO: Maybe can be more specific. Up for discussion. +OperatorType: TypeAlias = Callable[..., Any] SteppersOperatorsType: TypeAlias = tuple[tuple[OperatorType, ...], ...] RodType: TypeAlias = Type["RodBase"] SystemCollectionType: TypeAlias = "SystemCollectionProtocol" -AllowedContactType: TypeAlias = SystemType | Type["SurfaceBase"] +AllowedContactType: TypeAlias = ( + SystemType | Type["SurfaceBase"] +) # FIXME: SurfaceBase needs to be treated differently # Indexing types # TODO: Maybe just use slice?? From c1217b4cdd5b295ea613ae229c0b70044319331e Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Thu, 13 Jun 2024 21:40:55 +0900 Subject: [PATCH 066/134] remove deprecated functions to avoid unnecessary type hinting --- elastica/joint.py | 159 +------------------------ tests/test_joint.py | 284 -------------------------------------------- 2 files changed, 5 insertions(+), 438 deletions(-) diff --git a/elastica/joint.py b/elastica/joint.py index 260df2724..eb3fd5d18 100644 --- a/elastica/joint.py +++ b/elastica/joint.py @@ -1,6 +1,6 @@ __doc__ = """ Module containing joint classes to connect multiple rods together. """ -from typing import Any, NoReturn, Optional +from typing import Optional from elastica._rotations import _inv_rotate from elastica.typing import SystemType, RodType, ConnectionIndex @@ -59,11 +59,11 @@ def apply_forces( ---------- system_one : SystemType Rod or rigid-body object - index_one : int | np.ndarray + index_one : ConnectionIndex Index of first rod for joint. system_two : SystemType Rod or rigid-body object - index_two : int | np.ndarray + index_two : ConnectionIndex Index of second rod for joint. Returns @@ -104,11 +104,11 @@ def apply_torques( ---------- system_one : SystemType Rod or rigid-body object - index_one : int | np.ndarray + index_one : ConnectionIndex Index of first rod for joint. system_two : SystemType Rod or rigid-body object - index_two : int | np.ndarray + index_two : ConnectionIndex Index of second rod for joint. Returns @@ -386,155 +386,6 @@ def get_relative_rotation_two_systems( ) -# everything below this comment should be removed beyond v0.4.0 -def _dot_product(a: Any, b: Any) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._dot_product()\n" - "instead for find the dot product between a and b." - ) - - -def _norm(a: Any) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._norm()\n" - "instead for finding the norm of a." - ) - - -def _clip(x: Any, low: Any, high: Any) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._clip()\n" - "instead for clipping x." - ) - - -def _out_of_bounds(x: Any, low: Any, high: Any) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._out_of_bounds()\n" - "instead for checking if x is out of bounds." - ) - - -def _find_min_dist(x1: Any, e1: Any, x2: Any, e2: Any) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._find_min_dist()\n" - "instead for finding minimum distance between contact points." - ) - - -def _calculate_contact_forces_rod_rigid_body( - x_collection_rod: Any, - edge_collection_rod: Any, - x_cylinder_center: Any, - x_cylinder_tip: Any, - edge_cylinder: Any, - radii_sum: Any, - length_sum: Any, - internal_forces_rod: Any, - external_forces_rod: Any, - external_forces_cylinder: Any, - external_torques_cylinder: Any, - cylinder_director_collection: Any, - velocity_rod: Any, - velocity_cylinder: Any, - contact_k: Any, - contact_nu: Any, - velocity_damping_coefficient: Any, - friction_coefficient: Any, -) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. Please use\n" - "elastica._contact_functions._calculate_contact_forces_rod_cylinder()\n" - "instead for calculating rod cylinder contact forces." - ) - - -def _calculate_contact_forces_rod_rod( - x_collection_rod_one: Any, - radius_rod_one: Any, - length_rod_one: Any, - tangent_rod_one: Any, - velocity_rod_one: Any, - internal_forces_rod_one: Any, - external_forces_rod_one: Any, - x_collection_rod_two: Any, - radius_rod_two: Any, - length_rod_two: Any, - tangent_rod_two: Any, - velocity_rod_two: Any, - internal_forces_rod_two: Any, - external_forces_rod_two: Any, - contact_k: Any, - contact_nu: Any, -) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. Please use\n" - "elastica._contact_functions._calculate_contact_forces_rod_rod()\n" - "instead for calculating rod rod contact forces." - ) - - -def _calculate_contact_forces_self_rod( - x_collection_rod: Any, - radius_rod: Any, - length_rod: Any, - tangent_rod: Any, - velocity_rod: Any, - external_forces_rod: Any, - contact_k: Any, - contact_nu: Any, -) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. Please use\n" - "elastica._contact_functions._calculate_contact_forces_self_rod()\n" - "instead for calculating rod self-contact forces." - ) - - -def _aabbs_not_intersecting(aabb_one: Any, aabb_two: Any) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._aabbs_not_intersecting()\n" - "instead for checking aabbs intersection." - ) - - -def _prune_using_aabbs_rod_rigid_body( - rod_one_position_collection: Any, - rod_one_radius_collection: Any, - rod_one_length_collection: Any, - cylinder_position: Any, - cylinder_director: Any, - cylinder_radius: Any, - cylinder_length: Any, -) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._prune_using_aabbs_rod_cylinder()\n" - "instead for checking rod cylinder intersection." - ) - - -def _prune_using_aabbs_rod_rod( - rod_one_position_collection: Any, - rod_one_radius_collection: Any, - rod_one_length_collection: Any, - rod_two_position_collection: Any, - rod_two_radius_collection: Any, - rod_two_length_collection: Any, -) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._prune_using_aabbs_rod_rod()\n" - "instead for checking rod rod intersection." - ) - - class ExternalContact(FreeJoint): """ This class is for applying contact forces between rod-cylinder and rod-rod. diff --git a/tests/test_joint.py b/tests/test_joint.py index 373cd5f2d..40cbd7c27 100644 --- a/tests/test_joint.py +++ b/tests/test_joint.py @@ -376,119 +376,8 @@ def test_fixedjoint(rest_euler_angle): ) -from elastica.joint import ( - _dot_product, - _norm, - _clip, - _out_of_bounds, - _find_min_dist, - _aabbs_not_intersecting, -) - - -@pytest.mark.parametrize("ndim", [2, 3, 5, 10, 20]) -def test_dot_product_error_message(ndim): - vector1 = np.random.randn(ndim) - vector2 = np.random.randn(ndim) - error_message = ( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._dot_product()\n" - "instead for find the dot product between a and b." - ) - with pytest.raises(NotImplementedError) as error_info: - dot_product = _dot_product(vector1, vector2) - assert error_info.value.args[0] == error_message - - -@pytest.mark.parametrize("ndim", [2, 3, 5, 10, 20]) -def test_norm_error_message(ndim): - vec1 = np.random.randn(ndim) - error_message = ( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._norm()\n" - "instead for finding the norm of a." - ) - with pytest.raises(NotImplementedError) as error_info: - norm = _norm(vec1) - assert error_info.value.args[0] == error_message - - -@pytest.mark.parametrize( - "x, result", - [(0.5, 1), (1.5, 1.5), (2.5, 2)], -) -def test_clip_error_message(x, result): - low = 1.0 - high = 2.0 - error_message = ( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._clip()\n" - "instead for clipping x." - ) - with pytest.raises(NotImplementedError) as error_info: - _clip(x, low, high) - assert error_info.value.args[0] == error_message - - -@pytest.mark.parametrize( - "x, result", - [(0.5, 1), (1.5, 1.5), (2.5, 2)], -) -def test_out_of_bounds_error_message(x, result): - low = 1.0 - high = 2.0 - error_message = ( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._out_of_bounds()\n" - "instead for checking if x is out of bounds." - ) - with pytest.raises(NotImplementedError) as error_info: - _out_of_bounds(x, low, high) - assert error_info.value.args[0] == error_message - - -def test_find_min_dist_error_message(): - x1 = np.array([0, 0, 0]) - e1 = np.array([1, 1, 1]) - x2 = np.array([0, 1, 0]) - e2 = np.array([1, 0, 1]) - - error_message = ( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._find_min_dist()\n" - "instead for finding minimum distance between contact points." - ) - with pytest.raises(NotImplementedError) as error_info: - ( - min_dist_vec, - contact_point_of_system2, - contact_point_of_system1, - ) = _find_min_dist(x1, e1, x2, e2) - assert error_info.value.args[0] == error_message - - -def test_aabbs_not_intersecting_error_message(): - aabb_one = np.array([[0, 0], [0, 0], [0, 0]]) - aabb_two = np.array([[0, 0], [0, 0], [0, 0]]) - error_message = ( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._aabbs_not_intersecting()\n" - "instead for checking aabbs intersection." - ) - with pytest.raises(NotImplementedError) as error_info: - _aabbs_not_intersecting(aabb_one, aabb_two) - assert error_info.value.args[0] == error_message - - from elastica.rod import RodBase from elastica.rigidbody import RigidBodyBase -from elastica.joint import ( - _prune_using_aabbs_rod_rigid_body, - _prune_using_aabbs_rod_rod, - _calculate_contact_forces_rod_rigid_body, - _calculate_contact_forces_rod_rod, - _calculate_contact_forces_self_rod, -) def mock_rod_init(self): @@ -532,179 +421,6 @@ def mock_rigid_body_init(self): ) -def test_prune_using_aabbs_rod_rigid_body_error_message(): - rod = MockRod() - cylinder = MockRigidBody() - error_message = ( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._prune_using_aabbs_rod_cylinder()\n" - "instead for checking rod cylinder intersection." - ) - with pytest.raises(NotImplementedError) as error_info: - _prune_using_aabbs_rod_rigid_body( - rod.position_collection, - rod.radius, - rod.lengths, - cylinder.position_collection, - cylinder.director_collection, - cylinder.radius, - cylinder.length, - ) - assert error_info.value.args[0] == error_message - - -def test_prune_using_aabbs_rod_rod_error_message(): - rod_one = MockRod() - rod_two = MockRod() - error_message = ( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._prune_using_aabbs_rod_rod()\n" - "instead for checking rod rod intersection." - ) - with pytest.raises(NotImplementedError) as error_info: - _prune_using_aabbs_rod_rod( - rod_one.position_collection, - rod_one.radius, - rod_one.lengths, - rod_two.position_collection, - rod_two.radius, - rod_two.lengths, - ) - assert error_info.value.args[0] == error_message - - -def test_calculate_contact_forces_rod_rigid_body_error_message(): - - "initializing rod parameters" - rod = MockRod() - rod_element_position = 0.5 * ( - rod.position_collection[..., 1:] + rod.position_collection[..., :-1] - ) - - "initializing cylinder parameters" - cylinder = MockRigidBody() - x_cyl = ( - cylinder.position_collection[..., 0] - - 0.5 * cylinder.length * cylinder.director_collection[2, :, 0] - ) - - "initializing constants" - """ - Setting contact_k = 1 and other parameters to 0, - so the net forces becomes a function of contact forces only. - """ - k = 1.0 - nu = 0 - velocity_damping_coefficient = 0 - friction_coefficient = 0 - - error_message = ( - "This function is removed in v0.3.2. Please use\n" - "elastica._contact_functions._calculate_contact_forces_rod_cylinder()\n" - "instead for calculating rod cylinder contact forces." - ) - - "Function call" - with pytest.raises(NotImplementedError) as error_info: - _calculate_contact_forces_rod_rigid_body( - rod_element_position, - rod.lengths * rod.tangents, - cylinder.position_collection[..., 0], - x_cyl, - cylinder.length * cylinder.director_collection[2, :, 0], - rod.radius + cylinder.radius, - rod.lengths + cylinder.length, - rod.internal_forces, - rod.external_forces, - cylinder.external_forces, - cylinder.external_torques, - cylinder.director_collection[:, :, 0], - rod.velocity_collection, - cylinder.velocity_collection, - k, - nu, - velocity_damping_coefficient, - friction_coefficient, - ) - assert error_info.value.args[0] == error_message - - -def test_calculate_contact_forces_rod_rod_error_message(): - - rod_one = MockRod() - rod_two = MockRod() - """Placing rod two such that its first element just touches the last element of rod one.""" - rod_two.position_collection = np.array([[4, 5, 6], [0, 0, 0], [0, 0, 0]]) - - "initializing constants" - """ - Setting contact_k = 1 and nu to 0, - so the net forces becomes a function of contact forces only. - """ - k = 1.0 - nu = 0.0 - - error_message = ( - "This function is removed in v0.3.2. Please use\n" - "elastica._contact_functions._calculate_contact_forces_rod_rod()\n" - "instead for calculating rod rod contact forces." - ) - - "Function call" - with pytest.raises(NotImplementedError) as error_info: - _calculate_contact_forces_rod_rod( - rod_one.position_collection[..., :-1], - rod_one.radius, - rod_one.lengths, - rod_one.tangents, - rod_one.velocity_collection, - rod_one.internal_forces, - rod_one.external_forces, - rod_two.position_collection[..., :-1], - rod_two.radius, - rod_two.lengths, - rod_two.tangents, - rod_two.velocity_collection, - rod_two.internal_forces, - rod_two.external_forces, - k, - nu, - ) - assert error_info.value.args[0] == error_message - - -def test_calculate_contact_forces_self_rod_error_message(): - "Function to test the calculate contact forces self rod function" - - "Testing function with handcrafted/calculated values" - - rod = MockRod() - - "initializing constants" - k = 1.0 - nu = 1.0 - - error_message = ( - "This function is removed in v0.3.2. Please use\n" - "elastica._contact_functions._calculate_contact_forces_self_rod()\n" - "instead for calculating rod self-contact forces." - ) - - "Function call" - with pytest.raises(NotImplementedError) as error_info: - _calculate_contact_forces_self_rod( - rod.position_collection[..., :-1], - rod.radius, - rod.lengths, - rod.tangents, - rod.velocity_collection, - rod.external_forces, - k, - nu, - ) - assert error_info.value.args[0] == error_message - - class TestExternalContact: def test_external_contact_rod_rigid_body_with_collision_with_k_without_nu_and_friction( self, From 5d139d3b9f9b86a898dc7412573fb05199c3d758 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Thu, 13 Jun 2024 22:17:39 +0900 Subject: [PATCH 067/134] exclude experimental work from type-hinting --- pyproject.toml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 0816b9caa..f665dd07e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -142,6 +142,11 @@ warn_unreachable = false warn_unused_configs = true warn_unused_ignores = false +exclude = [ + "elastica/systems/analytical.py", + "elastica/experimental/*", +] + [tool.coverage.report] # Regexes for lines to exclude from consideration exclude_lines = [ From 3c7bc5aad47d326fdfe8b45e7b6d4d7f11e56968 Mon Sep 17 00:00:00 2001 From: Songyuan Cui Date: Tue, 7 May 2024 16:02:46 -0500 Subject: [PATCH 068/134] chore: internal _typing.py for typedefs for rebase --- elastica/rigidbody/__init__.py | 2 ++ elastica/rigidbody/_typing.py | 7 +++++++ 2 files changed, 9 insertions(+) create mode 100644 elastica/rigidbody/_typing.py diff --git a/elastica/rigidbody/__init__.py b/elastica/rigidbody/__init__.py index 55c6b9f69..9d3d4f979 100644 --- a/elastica/rigidbody/__init__.py +++ b/elastica/rigidbody/__init__.py @@ -2,3 +2,5 @@ from .cylinder import Cylinder from .sphere import Sphere from .mesh_rigid_body import MeshRigidBody + +from ._typing import float_t, f_arr_t, int_t diff --git a/elastica/rigidbody/_typing.py b/elastica/rigidbody/_typing.py new file mode 100644 index 000000000..941e1d411 --- /dev/null +++ b/elastica/rigidbody/_typing.py @@ -0,0 +1,7 @@ +import numpy as np +from numpy.typing import NDArray +from typing import TypeAlias + +float_t: TypeAlias = np.floating +f_arr_t: TypeAlias = NDArray[float_t] +int_t: TypeAlias = int From cb736c0872d9be00a2f4eac2369049181396359d Mon Sep 17 00:00:00 2001 From: Songyuan Cui Date: Tue, 7 May 2024 16:28:26 -0500 Subject: [PATCH 069/134] chore: typing for rigid_body.py --- elastica/rigidbody/rigid_body.py | 37 +++++++++++++++++--------------- 1 file changed, 20 insertions(+), 17 deletions(-) diff --git a/elastica/rigidbody/rigid_body.py b/elastica/rigidbody/rigid_body.py index e87392744..26edb014b 100644 --- a/elastica/rigidbody/rigid_body.py +++ b/elastica/rigidbody/rigid_body.py @@ -3,6 +3,9 @@ import numpy as np from abc import ABC from elastica._linalg import _batch_matvec, _batch_cross +from ._typing import f_arr_t, float_t + +from typing import Any class RigidBodyBase(ABC): @@ -17,22 +20,22 @@ class RigidBodyBase(ABC): REQUISITE_MODULES = [] - def __init__(self): + def __init__(self) -> None: - self.position_collection = NotImplementedError - self.velocity_collection = NotImplementedError - self.acceleration_collection = NotImplementedError - self.omega_collection = NotImplementedError - self.alpha_collection = NotImplementedError - self.director_collection = NotImplementedError + self.position_collection: f_arr_t + self.velocity_collection: f_arr_t + self.acceleration_collection: f_arr_t + self.omega_collection: f_arr_t + self.alpha_collection: f_arr_t + self.director_collection: f_arr_t - self.external_forces = NotImplementedError - self.external_torques = NotImplementedError + self.external_forces: f_arr_t + self.external_torques: f_arr_t - self.mass = NotImplementedError + self.mass: f_arr_t - self.mass_second_moment_of_inertia = NotImplementedError - self.inv_mass_second_moment_of_inertia = NotImplementedError + self.mass_second_moment_of_inertia: f_arr_t + self.inv_mass_second_moment_of_inertia: f_arr_t # @abstractmethod # # def update_accelerations(self): @@ -52,7 +55,7 @@ def __init__(self): # """ # pass - def update_accelerations(self, time): + def update_accelerations(self, time: float_t) -> None: np.copyto( self.acceleration_collection, (self.external_forces) / self.mass, @@ -74,18 +77,18 @@ def update_accelerations(self, time): ), ) - def zeroed_out_external_forces_and_torques(self, time): + def zeroed_out_external_forces_and_torques(self, time: float_t) -> None: # Reset forces and torques self.external_forces *= 0.0 self.external_torques *= 0.0 - def compute_position_center_of_mass(self): + def compute_position_center_of_mass(self) -> f_arr_t: """ Return positional center of mass """ return self.position_collection[..., 0].copy() - def compute_translational_energy(self): + def compute_translational_energy(self) -> Any: """ Return translational energy """ @@ -97,7 +100,7 @@ def compute_translational_energy(self): ) ) - def compute_rotational_energy(self): + def compute_rotational_energy(self) -> Any: """ Return rotational energy """ From ee7cc55df862c9bdec333578c403a9293c0dad40 Mon Sep 17 00:00:00 2001 From: Songyuan Cui Date: Tue, 7 May 2024 16:29:25 -0500 Subject: [PATCH 070/134] refactor: change J_omega to lower case --- elastica/rigidbody/rigid_body.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/elastica/rigidbody/rigid_body.py b/elastica/rigidbody/rigid_body.py index 26edb014b..a9737ca49 100644 --- a/elastica/rigidbody/rigid_body.py +++ b/elastica/rigidbody/rigid_body.py @@ -62,12 +62,12 @@ def update_accelerations(self, time: float_t) -> None: ) # I apply common sub expression elimination here, as J w - J_omega = _batch_matvec( + j_omega = _batch_matvec( self.mass_second_moment_of_inertia, self.omega_collection ) # (J \omega_L ) x \omega_L - lagrangian_transport = _batch_cross(J_omega, self.omega_collection) + lagrangian_transport = _batch_cross(j_omega, self.omega_collection) np.copyto( self.alpha_collection, @@ -104,7 +104,7 @@ def compute_rotational_energy(self) -> Any: """ Return rotational energy """ - J_omega = np.einsum( + j_omega = np.einsum( "ijk,jk->ik", self.mass_second_moment_of_inertia, self.omega_collection ) - return 0.5 * np.einsum("ik,ik->k", self.omega_collection, J_omega).sum() + return 0.5 * np.einsum("ik,ik->k", self.omega_collection, j_omega).sum() From b2e6ebf92681f81cc2f5ec7c9d8317428dc80a9f Mon Sep 17 00:00:00 2001 From: Songyuan Cui Date: Tue, 7 May 2024 16:41:13 -0500 Subject: [PATCH 071/134] chore: typing for cylinder.py --- elastica/rigidbody/cylinder.py | 89 ++++++++++++++++++++++++++-------- 1 file changed, 70 insertions(+), 19 deletions(-) diff --git a/elastica/rigidbody/cylinder.py b/elastica/rigidbody/cylinder.py index 58bdaa9e8..61d71d10e 100644 --- a/elastica/rigidbody/cylinder.py +++ b/elastica/rigidbody/cylinder.py @@ -6,9 +6,19 @@ from elastica.utils import MaxDimension from elastica.rigidbody.rigid_body import RigidBodyBase +from ._typing import f_arr_t, float_t, int_t + class Cylinder(RigidBodyBase): - def __init__(self, start, direction, normal, base_length, base_radius, density): + def __init__( + self, + start: f_arr_t, + direction: f_arr_t, + normal: f_arr_t, + base_length: float_t, + base_radius: float_t, + density: float_t, + ) -> None: """ Rigid body cylinder initializer. @@ -21,43 +31,80 @@ def __init__(self, start, direction, normal, base_length, base_radius, density): base_radius density """ - # rigid body does not have elements it only have one node. We are setting n_elems to - # zero for only make code to work. _bootstrap_from_data requires n_elems to be defined - self.n_elems = 1 - normal = normal.reshape(3, 1) - tangents = direction.reshape(3, 1) + def _check_array_size( + to_check: f_arr_t, name: str, expected: int_t = 3 + ) -> None: + array_size = to_check.size + assert array_size == expected, ( + f"Invalid size of '{name}'. " + f"Expected: {expected}, but got: {array_size}" + ) + + def _check_lower_bound( + to_check: float_t, name: str, lower_bound: float_t = np.float64(0.0) + ) -> None: + assert ( + to_check > lower_bound + ), f"Value for '{name}' ({to_check}) must be at lease {lower_bound}. " + + _check_array_size(start, "start") + _check_array_size(direction, "direction") + _check_array_size(normal, "normal") + + _check_lower_bound(base_length, "base_length") + _check_lower_bound(base_radius, "base_radius") + _check_lower_bound(density, "density") + + super().__init__() + + # rigid body does not have elements it only has one node. We are setting n_elems to + # zero for only make code to work. _bootstrap_from_data requires n_elems to be define + self.n_elem: int_t = 1 + + normal = normal.reshape((3, 1)) + tangents = direction.reshape((3, 1)) binormal = _batch_cross(tangents, normal) self.radius = base_radius self.length = base_length self.density = density + + dim: int_t = MaxDimension.value() + # This is for a rigid body cylinder self.volume = np.pi * base_radius * base_radius * base_length - self.mass = np.array([self.volume * self.density]) + self.mass = np.array([self.volume * self.density], dtype=np.float64) # Second moment of inertia - A0 = np.pi * base_radius * base_radius - I0_1 = A0 * A0 / (4.0 * np.pi) - I0_2 = I0_1 - I0_3 = 2.0 * I0_2 - I0 = np.array([I0_1, I0_2, I0_3]) + area = np.pi * base_radius * base_radius + smoa_span_1 = area * area / (4.0 * np.pi) + smoa_span_2 = smoa_span_1 + smoa_axial = 2.0 * smoa_span_1 + smoa = np.array([smoa_span_1, smoa_span_2, smoa_axial]) + + # Allocate properties + self.position_collection = np.zeros((dim, 1), dtype=np.float64) + self.velocity_collection = np.zeros((dim, 1), dtype=np.float64) + self.acceleration_collection = np.zeros((dim, 1), dtype=np.float64) + self.omega_collection = np.zeros((dim, 1), dtype=np.float64) + self.alpha_collection = np.zeros((dim, 1), dtype=np.float64) + self.director_collection = np.zeros((dim, dim, 1), dtype=np.float64) + + self.external_forces = np.zeros((dim, 1), dtype=np.float64) + self.external_torques = np.zeros((dim, 1), dtype=np.float64) # Mass second moment of inertia for disk cross-section - mass_second_moment_of_inertia = np.zeros( - (MaxDimension.value(), MaxDimension.value()), np.float64 - ) - np.fill_diagonal(mass_second_moment_of_inertia, I0 * density * base_length) + mass_second_moment_of_inertia = np.diag(smoa * density * base_length) self.mass_second_moment_of_inertia = mass_second_moment_of_inertia.reshape( - MaxDimension.value(), MaxDimension.value(), 1 + (dim, dim, 1) ) self.inv_mass_second_moment_of_inertia = np.linalg.inv( mass_second_moment_of_inertia - ).reshape(MaxDimension.value(), MaxDimension.value(), 1) + ).reshape((dim, dim, 1)) # position is at the center - self.position_collection = np.zeros((MaxDimension.value(), 1)) self.position_collection[:] = ( start.reshape(3, 1) + direction.reshape(3, 1) * base_length / 2 ) @@ -80,3 +127,7 @@ def __init__(self, start, direction, normal, base_length, base_radius, density): self.external_torques = np.zeros((MaxDimension.value())).reshape( MaxDimension.value(), 1 ) + # Assemble directors + self.director_collection[0, :] = normal + self.director_collection[1, :] = binormal + self.director_collection[2, :] = tangents From 0febaa059df1ff3878a60080b54c8e6d36ed4a35 Mon Sep 17 00:00:00 2001 From: Songyuan Cui Date: Tue, 7 May 2024 16:50:03 -0500 Subject: [PATCH 072/134] chore: remove commented functions --- elastica/rigidbody/rigid_body.py | 18 ------------------ 1 file changed, 18 deletions(-) diff --git a/elastica/rigidbody/rigid_body.py b/elastica/rigidbody/rigid_body.py index a9737ca49..9ecf75630 100644 --- a/elastica/rigidbody/rigid_body.py +++ b/elastica/rigidbody/rigid_body.py @@ -37,24 +37,6 @@ def __init__(self) -> None: self.mass_second_moment_of_inertia: f_arr_t self.inv_mass_second_moment_of_inertia: f_arr_t - # @abstractmethod - # # def update_accelerations(self): - # # pass - - # def _compute_internal_forces_and_torques(self): - # """ - # This function here is only for integrator to work properly. We do not need - # internal forces and torques at all. - # Parameters - # ---------- - # time - # - # Returns - # ------- - # - # """ - # pass - def update_accelerations(self, time: float_t) -> None: np.copyto( self.acceleration_collection, From e1940881e3f45ef0b405c69ba8a3d1da8b02444b Mon Sep 17 00:00:00 2001 From: Songyuan Cui Date: Tue, 7 May 2024 16:58:46 -0500 Subject: [PATCH 073/134] chore: typing for sphere.py --- elastica/rigidbody/sphere.py | 58 ++++++++++++++++++++---------------- 1 file changed, 32 insertions(+), 26 deletions(-) diff --git a/elastica/rigidbody/sphere.py b/elastica/rigidbody/sphere.py index 181dd18a5..92c68ac96 100644 --- a/elastica/rigidbody/sphere.py +++ b/elastica/rigidbody/sphere.py @@ -6,9 +6,11 @@ from elastica.utils import MaxDimension from elastica.rigidbody.rigid_body import RigidBodyBase +from ._typing import float_t, f_arr_t, int_t + class Sphere(RigidBodyBase): - def __init__(self, center, base_radius, density): + def __init__(self, center: f_arr_t, base_radius: float_t, density: float_t) -> None: """ Rigid body sphere initializer. @@ -18,55 +20,59 @@ def __init__(self, center, base_radius, density): base_radius density """ + + super().__init__() + # rigid body does not have elements it only have one node. We are setting n_elems to # zero for only make code to work. _bootstrap_from_data requires n_elems to be defined - self.n_elems = 1 + self.n_elems: int_t = 1 + + dim: int_t = MaxDimension.value() + + assert ( + center.size == dim + ), f"center must be of size {dim}, but was {center.size}" + assert base_radius > 0.0, "base_radius must be positive" + assert density > 0.0, "density must be positive" self.radius = base_radius self.density = density self.length = 2 * base_radius # This is for a rigid body cylinder self.volume = 4.0 / 3.0 * np.pi * base_radius**3 - self.mass = np.array([self.volume * self.density]) - normal = np.array([1.0, 0.0, 0.0]).reshape(3, 1) - tangents = np.array([0.0, 0.0, 1.0]).reshape(3, 1) + self.mass = np.array([self.volume * self.density], dtype=np.float64) + normal = np.array([1.0, 0.0, 0.0], dtype=np.float64).reshape(dim, 1) + tangents = np.array([0.0, 0.0, 1.0], dtype=np.float64).reshape(dim, 1) binormal = _batch_cross(tangents, normal) # Mass second moment of inertia for disk cross-section - mass_second_moment_of_inertia = np.zeros( - (MaxDimension.value(), MaxDimension.value()), np.float64 - ) + mass_second_moment_of_inertia = np.zeros((dim, dim), dtype=np.float64) np.fill_diagonal( mass_second_moment_of_inertia, 2.0 / 5.0 * self.mass * self.radius**2 ) self.mass_second_moment_of_inertia = mass_second_moment_of_inertia.reshape( - MaxDimension.value(), MaxDimension.value(), 1 + (dim, dim, 1) ) self.inv_mass_second_moment_of_inertia = np.linalg.inv( mass_second_moment_of_inertia - ).reshape(MaxDimension.value(), MaxDimension.value(), 1) + ).reshape((dim, dim, 1)) + + # Allocate properties + self.position_collection = np.zeros((dim, 1), dtype=np.float64) + self.velocity_collection = np.zeros((dim, 1), dtype=np.float64) + self.acceleration_collection = np.zeros((dim, 1), dtype=np.float64) + self.omega_collection = np.zeros((dim, 1), dtype=np.float64) + self.alpha_collection = np.zeros((dim, 1), dtype=np.float64) + self.director_collection = np.zeros((dim, dim, 1), dtype=np.float64) + + self.external_forces = np.zeros((dim, 1), dtype=np.float64) + self.external_torques = np.zeros((dim, 1), dtype=np.float64) # position is at the center - self.position_collection = np.zeros((MaxDimension.value(), 1)) self.position_collection[:, 0] = center - self.velocity_collection = np.zeros((MaxDimension.value(), 1)) - self.omega_collection = np.zeros((MaxDimension.value(), 1)) - self.acceleration_collection = np.zeros((MaxDimension.value(), 1)) - self.alpha_collection = np.zeros((MaxDimension.value(), 1)) - - self.director_collection = np.zeros( - (MaxDimension.value(), MaxDimension.value(), 1) - ) self.director_collection[0, ...] = normal self.director_collection[1, ...] = binormal self.director_collection[2, ...] = tangents - - self.external_forces = np.zeros((MaxDimension.value())).reshape( - MaxDimension.value(), 1 - ) - self.external_torques = np.zeros((MaxDimension.value())).reshape( - MaxDimension.value(), 1 - ) From 8f0ad693378b3404c6c16e7264a741d1f0422108 Mon Sep 17 00:00:00 2001 From: Songyuan Cui Date: Fri, 10 May 2024 09:35:22 -0500 Subject: [PATCH 074/134] refactor: remove unused _bootstrap functions and state classes --- elastica/rigidbody/data_structures.py | 435 +++++++++++++------------- 1 file changed, 216 insertions(+), 219 deletions(-) diff --git a/elastica/rigidbody/data_structures.py b/elastica/rigidbody/data_structures.py index e68f7c295..bbca69886 100644 --- a/elastica/rigidbody/data_structures.py +++ b/elastica/rigidbody/data_structures.py @@ -1,8 +1,5 @@ __doc__ = "Data structure wrapper for rod components" -import numpy as np - -from elastica._rotations import _get_rotation_matrix from elastica.rod.data_structures import _RodSymplecticStepperMixin """ @@ -55,79 +52,79 @@ def update_internal_forces_and_torques(self, *args, **kwargs) -> None: pass -def _bootstrap_from_data(stepper_type: str, n_elems: int, vector_states, matrix_states): - """Returns states wrapping numpy arrays based on the time-stepping algorithm - - Convenience method that takes in rod internal (raw np.ndarray) data, create views - (references) from it, and outputs State classes that are used in the time-stepping - algorithm. This means that modifying the state modifies the internal data! - - Parameters - ---------- - stepper_type : str (likely to change in future), representing stepper type - Allowed parameters are ['explicit', 'symplectic'] - n_elems : int, number of rod elements - vector_states : np.ndarray of shape (dim, *) with the following structure - `vector_states` = [`position`,`velocity`,`omega`,`acceleration`,`angular acceleration`] - `n_nodes = n_elems = 1` - `position = 0 -> n_nodes , size = n_nodes` - `velocity = n_nodes -> 2 * n_nodes, size = n_nodes` - `omega = 2 * n_nodes -> 2 * n_nodes + nelem, size = nelem` - `acceleration = 2 * n_nodes + nelem -> 3 * n_nodes + nelem, size = n_nodes` - `angular acceleration = 3 * n_nodes + nelem -> 3 * n_nodes + 2 * nelem, size = n_elems` - matrix_states : np.ndarray of shape (dim, dim, n_elems) containing the directors - - Returns - ------- - output : tuple of len 8 containing - (state, derivative_state, position, directors, velocity, omega, acceleration, alpha) - derivative_state carries rate information - - """ - n_nodes = n_elems - position = np.ndarray.view(vector_states[..., :n_nodes]) - directors = np.ndarray.view(matrix_states) - v_w_dvdt_dwdt = np.ndarray.view(vector_states[..., n_nodes:]) - output = () - # TODO: - # 11/01/2020: Extend Rigid body data structures for Explicit steppers - # 12/20/2021: Future work! If explicit stepper is gonna be dropped, remove. - # if stepper_type == "explicit": - # v_w_states = np.ndarray.view(vector_states[..., n_nodes : n_nodes + 1]) - # output += ( - # _State(n_elems, position, directors, v_w_states), - # _DerivativeState(n_elems, v_w_dvdt_dwdt), - # ) - # elif stepper_type == "symplectic": - # output += ( - # _KinematicState(n_elems, position, directors), - # _DynamicState(n_elems, v_w_dvdt_dwdt), - # ) - # else: - # return - output += ( - _KinematicState(n_elems, position, directors), - _DynamicState(n_elems, v_w_dvdt_dwdt), - ) - - n_velocity_end = n_nodes + n_nodes - velocity = np.ndarray.view(vector_states[..., n_nodes:n_velocity_end]) - - n_omega_end = n_velocity_end + n_elems - omega = np.ndarray.view(vector_states[..., n_velocity_end:n_omega_end]) - - n_acceleration_end = n_omega_end + n_nodes - acceleration = np.ndarray.view(vector_states[..., n_omega_end:n_acceleration_end]) - - n_alpha_end = n_acceleration_end + n_elems - alpha = np.ndarray.view(vector_states[..., n_acceleration_end:n_alpha_end]) - - return output + (position, directors, velocity, omega, acceleration, alpha) +# def _bootstrap_from_data(stepper_type: str, n_elems: int, vector_states, matrix_states): +# """Returns states wrapping numpy arrays based on the time-stepping algorithm +# +# Convenience method that takes in rod internal (raw np.ndarray) data, create views +# (references) from it, and outputs State classes that are used in the time-stepping +# algorithm. This means that modifying the state modifies the internal data! +# +# Parameters +# ---------- +# stepper_type : str (likely to change in future), representing stepper type +# Allowed parameters are ['explicit', 'symplectic'] +# n_elems : int, number of rod elements +# vector_states : np.ndarray of shape (dim, *) with the following structure +# `vector_states` = [`position`,`velocity`,`omega`,`acceleration`,`angular acceleration`] +# `n_nodes = n_elems = 1` +# `position = 0 -> n_nodes , size = n_nodes` +# `velocity = n_nodes -> 2 * n_nodes, size = n_nodes` +# `omega = 2 * n_nodes -> 2 * n_nodes + nelem, size = nelem` +# `acceleration = 2 * n_nodes + nelem -> 3 * n_nodes + nelem, size = n_nodes` +# `angular acceleration = 3 * n_nodes + nelem -> 3 * n_nodes + 2 * nelem, size = n_elems` +# matrix_states : np.ndarray of shape (dim, dim, n_elems) containing the directors +# +# Returns +# ------- +# output : tuple of len 8 containing +# (state, derivative_state, position, directors, velocity, omega, acceleration, alpha) +# derivative_state carries rate information +# +# """ +# n_nodes = n_elems +# position = np.ndarray.view(vector_states[..., :n_nodes]) +# directors = np.ndarray.view(matrix_states) +# v_w_dvdt_dwdt = np.ndarray.view(vector_states[..., n_nodes:]) +# output = () +# TODO: +# 11/01/2020: Extend Rigid body data structures for Explicit steppers +# 12/20/2021: Future work! If explicit stepper is gonna be dropped, remove. +# if stepper_type == "explicit": +# v_w_states = np.ndarray.view(vector_states[..., n_nodes : n_nodes + 1]) +# output += ( +# _State(n_elems, position, directors, v_w_states), +# _DerivativeState(n_elems, v_w_dvdt_dwdt), +# ) +# elif stepper_type == "symplectic": +# output += ( +# _KinematicState(n_elems, position, directors), +# _DynamicState(n_elems, v_w_dvdt_dwdt), +# ) +# else: +# return +# output += ( +# _KinematicState(n_elems, position, directors), +# _DynamicState(n_elems, v_w_dvdt_dwdt), +# ) +# +# n_velocity_end = n_nodes + n_nodes +# velocity = np.ndarray.view(vector_states[..., n_nodes:n_velocity_end]) +# +# n_omega_end = n_velocity_end + n_elems +# omega = np.ndarray.view(vector_states[..., n_velocity_end:n_omega_end]) +# +# n_acceleration_end = n_omega_end + n_nodes +# acceleration = np.ndarray.view(vector_states[..., n_omega_end:n_acceleration_end]) +# +# n_alpha_end = n_acceleration_end + n_elems +# alpha = np.ndarray.view(vector_states[..., n_acceleration_end:n_alpha_end]) +# +# return output + (position, directors, velocity, omega, acceleration, alpha) -""" -Explicit stepper interface -""" +# """ +# Explicit stepper interface +# """ # TODO # 12/20/2021: If explicit stepper is gonna be dropped, remove. @@ -362,152 +359,152 @@ def _bootstrap_from_data(stepper_type: str, n_elems: int, vector_states, matrix_ # """ # return self.__rmul__(scalar) -""" -Symplectic stepper interface -""" +# """ +# Symplectic stepper interface +# """ # TODO: Maybe considerg removing. We no longer use bootstrap to initialize. # RigidBodySymplecticStepperMixin is now derived from RodSymplecticStepperMixin -class _KinematicState: - """State storing (x,Q) for symplectic steppers. - - Wraps data as state, with overloaded methods for symplectic steppers. - Allows for separating implementation of stepper from actual - addition/multiplication/other formulae used. - - Symplectic steppers rely only on in-place modifications to state and so - only these methods are provided. - """ - - def __init__(self, position_collection_view, director_collection_view): - """ - Parameters - ---------- - n_elems : int, number of rod elements - position_collection_view : view of positions (or) x - director_collection_view : view of directors (or) Q - """ - self.position_collection = position_collection_view - self.director_collection = director_collection_view - - def __iadd__(self, scaled_deriv_array): - """overloaded += operator - - The add for directors is customized to reflect Rodrigues' rotation - formula. - - Parameters - ---------- - scaled_deriv_array : np.ndarray containing dt * (v, ω), - as retured from _DynamicState's `kinematic_rates` method - - Returns - ------- - self : _KinematicState instance with inplace modified data - - Caveats - ------- - Takes a np.ndarray and not a _KinematicState object (as one expects). - This is done for efficiency reasons, see _DynamicState's `kinematic_rates` - method - """ - velocity_collection = scaled_deriv_array[0] - omega_collection = scaled_deriv_array[1] - # x += v*dt - self.position_collection += velocity_collection - # Devs : see `_State.__iadd__` for reasons why we do matmul here - np.einsum( - "ijk,jlk->ilk", - _get_rotation_matrix(1.0, omega_collection), - self.director_collection.copy(), - out=self.director_collection, - ) - return self - - -class _DynamicState: - """State storing (v,ω, dv/dt, dω/dt) for symplectic steppers. - - Wraps data as state, with overloaded methods for symplectic steppers. - Allows for separating implementation of stepper from actual - addition/multiplication/other formulae used. - - Symplectic steppers rely only on in-place modifications to state and so - only these methods are provided. - """ - - def __init__( - self, - v_w_collection, - dvdt_dwdt_collection, - velocity_collection, - omega_collection, - ): - """ - - Parameters - ---------- - n_elems : int, number of rod elements - rate_collection_view : np.ndarray containing (v, ω, dv/dt, dω/dt) - """ - super(_DynamicState, self).__init__() - # Limit at which (v, w) end - self.rate_collection = v_w_collection - self.dvdt_dwdt_collection = dvdt_dwdt_collection - self.velocity_collection = velocity_collection - self.omega_collection = omega_collection - - def __iadd__(self, scaled_second_deriv_array): - """overloaded += operator, updating dynamic_rates - - Parameters - ---------- - scaled_second_deriv_array : np.ndarray containing dt * (dvdt, dωdt), - as retured from _DynamicState's `dynamic_rates` method - - Returns - ------- - self : _DynamicState instance with inplace modified data - - Caveats - ------- - Takes a np.ndarray and not a _DynamicState object (as one expects). - This is done for efficiency reasons, see `dynamic_rates`. - """ - # Always goes in LHS : that means the update is on the rates alone - # (v,ω) += dt * (dv/dt, dω/dt) -> self.dynamic_rates - self.rate_collection += scaled_second_deriv_array - return self - - def kinematic_rates(self, time, prefac, *args, **kwargs): - """Yields kinematic rates to interact with _KinematicState - - Returns - ------- - v_and_omega : np.ndarray consisting of (v,ω) - - Caveats - ------- - Doesn't return a _KinematicState with (dt*v, dt*w) as members, - as one expects the _Kinematic __add__ operator to interact - with another _KinematicState. This is done for efficiency purposes. - """ - # RHS functino call, gives v,w so that - # Comes from kin_state -> (x,Q) += dt * (v,w) <- First part of dyn_state - return prefac * self.velocity_collection, prefac * self.omega_collection - - def dynamic_rates(self, time, prefac, *args, **kwargs): - """Yields dynamic rates to add to with _DynamicState - - Returns - ------- - acc_and_alpha : np.ndarray consisting of (dv/dt,dω/dt) - - Caveats - ------- - Doesn't return a _DynamicState with (dt*v, dt*w) as members, - as one expects the _Dynamic __add__ operator to interact - with another _DynamicState. This is done for efficiency purposes. - """ - return prefac * self.dvdt_dwdt_collection +# class _KinematicState: +# """State storing (x,Q) for symplectic steppers. +# +# Wraps data as state, with overloaded methods for symplectic steppers. +# Allows for separating implementation of stepper from actual +# addition/multiplication/other formulae used. +# +# Symplectic steppers rely only on in-place modifications to state and so +# only these methods are provided. +# """ +# +# def __init__(self, position_collection_view, director_collection_view): +# """ +# Parameters +# ---------- +# n_elems : int, number of rod elements +# position_collection_view : view of positions (or) x +# director_collection_view : view of directors (or) Q +# """ +# self.position_collection = position_collection_view +# self.director_collection = director_collection_view +# +# def __iadd__(self, scaled_deriv_array): +# """overloaded += operator +# +# The add for directors is customized to reflect Rodrigues' rotation +# formula. +# +# Parameters +# ---------- +# scaled_deriv_array : np.ndarray containing dt * (v, ω), +# as retured from _DynamicState's `kinematic_rates` method +# +# Returns +# ------- +# self : _KinematicState instance with inplace modified data +# +# Caveats +# ------- +# Takes a np.ndarray and not a _KinematicState object (as one expects). +# This is done for efficiency reasons, see _DynamicState's `kinematic_rates` +# method +# """ +# velocity_collection = scaled_deriv_array[0] +# omega_collection = scaled_deriv_array[1] +# # x += v*dt +# self.position_collection += velocity_collection +# # Devs : see `_State.__iadd__` for reasons why we do matmul here +# np.einsum( +# "ijk,jlk->ilk", +# _get_rotation_matrix(1.0, omega_collection), +# self.director_collection.copy(), +# out=self.director_collection, +# ) +# return self +# +# +# class _DynamicState: +# """State storing (v,ω, dv/dt, dω/dt) for symplectic steppers. +# +# Wraps data as state, with overloaded methods for symplectic steppers. +# Allows for separating implementation of stepper from actual +# addition/multiplication/other formulae used. +# +# Symplectic steppers rely only on in-place modifications to state and so +# only these methods are provided. +# """ +# +# def __init__( +# self, +# v_w_collection, +# dvdt_dwdt_collection, +# velocity_collection, +# omega_collection, +# ): +# """ +# +# Parameters +# ---------- +# n_elems : int, number of rod elements +# rate_collection_view : np.ndarray containing (v, ω, dv/dt, dω/dt) +# """ +# super(_DynamicState, self).__init__() +# # Limit at which (v, w) end +# self.rate_collection = v_w_collection +# self.dvdt_dwdt_collection = dvdt_dwdt_collection +# self.velocity_collection = velocity_collection +# self.omega_collection = omega_collection +# +# def __iadd__(self, scaled_second_deriv_array): +# """overloaded += operator, updating dynamic_rates +# +# Parameters +# ---------- +# scaled_second_deriv_array : np.ndarray containing dt * (dvdt, dωdt), +# as retured from _DynamicState's `dynamic_rates` method +# +# Returns +# ------- +# self : _DynamicState instance with inplace modified data +# +# Caveats +# ------- +# Takes a np.ndarray and not a _DynamicState object (as one expects). +# This is done for efficiency reasons, see `dynamic_rates`. +# """ +# # Always goes in LHS : that means the update is on the rates alone +# # (v,ω) += dt * (dv/dt, dω/dt) -> self.dynamic_rates +# self.rate_collection += scaled_second_deriv_array +# return self +# +# def kinematic_rates(self, time, prefac, *args, **kwargs): +# """Yields kinematic rates to interact with _KinematicState +# +# Returns +# ------- +# v_and_omega : np.ndarray consisting of (v,ω) +# +# Caveats +# ------- +# Doesn't return a _KinematicState with (dt*v, dt*w) as members, +# as one expects the _Kinematic __add__ operator to interact +# with another _KinematicState. This is done for efficiency purposes. +# """ +# # RHS functino call, gives v,w so that +# # Comes from kin_state -> (x,Q) += dt * (v,w) <- First part of dyn_state +# return prefac * self.velocity_collection, prefac * self.omega_collection +# +# def dynamic_rates(self, time, prefac, *args, **kwargs): +# """Yields dynamic rates to add to with _DynamicState +# +# Returns +# ------- +# acc_and_alpha : np.ndarray consisting of (dv/dt,dω/dt) +# +# Caveats +# ------- +# Doesn't return a _DynamicState with (dt*v, dt*w) as members, +# as one expects the _Dynamic __add__ operator to interact +# with another _DynamicState. This is done for efficiency purposes. +# """ +# return prefac * self.dvdt_dwdt_collection From c9d1e95c2b9b62aa4c9b2634ff1daec196a932f1 Mon Sep 17 00:00:00 2001 From: Songyuan Cui Date: Fri, 10 May 2024 09:39:18 -0500 Subject: [PATCH 075/134] chore: typing for rigid body data structure --- elastica/rigidbody/data_structures.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/elastica/rigidbody/data_structures.py b/elastica/rigidbody/data_structures.py index bbca69886..6cf64a6d6 100644 --- a/elastica/rigidbody/data_structures.py +++ b/elastica/rigidbody/data_structures.py @@ -1,6 +1,7 @@ __doc__ = "Data structure wrapper for rod components" from elastica.rod.data_structures import _RodSymplecticStepperMixin +from typing import Any """ # FIXME : Explicit Stepper doesn't work as States lose the @@ -48,7 +49,7 @@ def __init__(self) -> None: # dynamic rates needs to call update_accelerations and henc # is another function - def update_internal_forces_and_torques(self, *args, **kwargs) -> None: + def update_internal_forces_and_torques(self, *args: Any, **kwargs: Any) -> None: pass From 780ed6f70cc834fbfccf395d7d059fd8622fcecb Mon Sep 17 00:00:00 2001 From: Songyuan Cui Date: Fri, 14 Jun 2024 22:24:42 +0800 Subject: [PATCH 076/134] fix: fixed repetitive codes in cylinder.py caused by rebase merge --- elastica/rigidbody/cylinder.py | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/elastica/rigidbody/cylinder.py b/elastica/rigidbody/cylinder.py index 61d71d10e..079ce7d53 100644 --- a/elastica/rigidbody/cylinder.py +++ b/elastica/rigidbody/cylinder.py @@ -109,25 +109,6 @@ def _check_lower_bound( start.reshape(3, 1) + direction.reshape(3, 1) * base_length / 2 ) - self.velocity_collection = np.zeros((MaxDimension.value(), 1)) - self.omega_collection = np.zeros((MaxDimension.value(), 1)) - self.acceleration_collection = np.zeros((MaxDimension.value(), 1)) - self.alpha_collection = np.zeros((MaxDimension.value(), 1)) - - self.director_collection = np.zeros( - (MaxDimension.value(), MaxDimension.value(), 1) - ) self.director_collection[0, ...] = normal self.director_collection[1, ...] = binormal self.director_collection[2, ...] = tangents - - self.external_forces = np.zeros((MaxDimension.value())).reshape( - MaxDimension.value(), 1 - ) - self.external_torques = np.zeros((MaxDimension.value())).reshape( - MaxDimension.value(), 1 - ) - # Assemble directors - self.director_collection[0, :] = normal - self.director_collection[1, :] = binormal - self.director_collection[2, :] = tangents From 60d951fd2b6a6225f7cc9c7e1908c42d318499c7 Mon Sep 17 00:00:00 2001 From: Songyuan Cui Date: Fri, 14 Jun 2024 22:26:57 +0800 Subject: [PATCH 077/134] refactor: temporarily remove REQUISITE_MODULES --- elastica/rigidbody/rigid_body.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/elastica/rigidbody/rigid_body.py b/elastica/rigidbody/rigid_body.py index 9ecf75630..12b70a2d9 100644 --- a/elastica/rigidbody/rigid_body.py +++ b/elastica/rigidbody/rigid_body.py @@ -18,8 +18,6 @@ class RigidBodyBase(ABC): """ - REQUISITE_MODULES = [] - def __init__(self) -> None: self.position_collection: f_arr_t From 62c3e21880beef69382c9a1630f7ccbdb83f2809 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Fri, 14 Jun 2024 23:57:26 +0900 Subject: [PATCH 078/134] typing: rearange protocol for system --- elastica/memory_block/protocol.py | 12 +++++++ elastica/rod/protocol.py | 46 +++++++++++++++++++++++++ elastica/rod/rod_base.py | 7 ++-- elastica/systems/protocol.py | 56 +++++++++++++++++-------------- 4 files changed, 93 insertions(+), 28 deletions(-) create mode 100644 elastica/memory_block/protocol.py create mode 100644 elastica/rod/protocol.py diff --git a/elastica/memory_block/protocol.py b/elastica/memory_block/protocol.py new file mode 100644 index 000000000..5cf002462 --- /dev/null +++ b/elastica/memory_block/protocol.py @@ -0,0 +1,12 @@ +from typing import Protocol + +from elastica.rod.protocol import CosseratRodProtocol +from elastica.systems.protocol import SymplecticSystemProtocol + + +class BlockCosseratRodProtocol(CosseratRodProtocol, SymplecticSystemProtocol, Protocol): + pass + + +# class BlockRigidBodyProtocol(RigidBodyProtocol, SymplecticSystemProtocol, Protocol): +# pass diff --git a/elastica/rod/protocol.py b/elastica/rod/protocol.py new file mode 100644 index 000000000..ba83bb225 --- /dev/null +++ b/elastica/rod/protocol.py @@ -0,0 +1,46 @@ +from typing import Protocol + +import numpy as np +from numpy.typing import NDArray + +pass + +from elastica.systems.protocol import SystemProtocol + + +class _RodEnergy(Protocol): + def compute_bending_energy(self) -> NDArray[np.float64]: ... + + def compute_shear_energy(self) -> NDArray[np.float64]: ... + + +class CosseratRodProtocol(SystemProtocol, _RodEnergy, Protocol): + + mass: NDArray[np.floating] + volume: NDArray[np.floating] + radius: NDArray[np.floating] + tangents: NDArray[np.floating] + lengths: NDArray[np.floating] + rest_lengths: NDArray[np.floating] + rest_voronoi_lengths: NDArray[np.floating] + kappa: NDArray[np.floating] + sigma: NDArray[np.floating] + rest_kappa: NDArray[np.floating] + rest_sigma: NDArray[np.floating] + + internal_stress: NDArray[np.floating] + internal_couple: NDArray[np.floating] + dilatation: NDArray[np.floating] + dilatation_rate: NDArray[np.floating] + voronoi_dilatation: NDArray[np.floating] + + bend_matrix: NDArray[np.floating] + shear_matrix: NDArray[np.floating] + + mass_second_moment_of_inertia: NDArray[np.floating] + inv_mass_second_moment_of_inertia: NDArray[np.floating] + + ghost_voronoi_idx: NDArray[np.integer] + ghost_elems_idx: NDArray[np.integer] + + ring_rod_flag: bool diff --git a/elastica/rod/rod_base.py b/elastica/rod/rod_base.py index dab11221f..76e262f15 100644 --- a/elastica/rod/rod_base.py +++ b/elastica/rod/rod_base.py @@ -1,6 +1,6 @@ __doc__ = """Base class for rods""" -from typing import Any +from typing import Type import numpy as np from numpy.typing import NDArray @@ -15,7 +15,7 @@ class RodBase: """ - REQUISITE_MODULES: list[Any] = [] + REQUISITE_MODULES: list[Type] = [] def __init__(self) -> None: """ @@ -29,3 +29,6 @@ def __init__(self) -> None: self.alpha_collection: NDArray[np.floating] self.external_forces: NDArray[np.floating] self.external_torques: NDArray[np.floating] + + self.ghost_voronoi_idx: NDArray[np.integer] + self.ghost_elems_idx: NDArray[np.integer] diff --git a/elastica/systems/protocol.py b/elastica/systems/protocol.py index 5855e0bad..d49129743 100644 --- a/elastica/systems/protocol.py +++ b/elastica/systems/protocol.py @@ -2,7 +2,6 @@ from typing import Protocol, Type from elastica.typing import StateType, SystemType -from elastica.modules.protocol import ModuleProtocol from elastica.rod.data_structures import _KinematicState, _DynamicState @@ -10,12 +9,24 @@ from numpy.typing import NDArray -class SystemProtocol(Protocol): +class _SystemWithEnergy(Protocol): + def compute_translational_energy(self) -> NDArray[np.float64]: ... + + def compute_rotational_energy(self) -> NDArray[np.float64]: ... + + +class _SystemWithCenterOfMass(Protocol): + def compute_velocity_center_of_mass(self) -> NDArray[np.float64]: ... + + def compute_position_center_of_mass(self) -> NDArray[np.float64]: ... + + +class SystemProtocol(_SystemWithEnergy, _SystemWithCenterOfMass, Protocol): """ Protocol for all elastica system """ - REQUISITE_MODULES: list[Type[ModuleProtocol]] + REQUISITE_MODULES: list[Type] @property def n_nodes(self) -> int: ... @@ -23,31 +34,27 @@ def n_nodes(self) -> int: ... @property def n_elems(self) -> int: ... - @property - def position_collection(self) -> NDArray: ... + position_collection: NDArray[np.floating] - @property - def velocity_collection(self) -> NDArray: ... + velocity_collection: NDArray[np.floating] - @property - def director_collection(self) -> NDArray: ... + acceleration_collection: NDArray[np.floating] + director_collection: NDArray[np.floating] - @property - def acceleration_collection(self) -> NDArray: ... + omega_collection: NDArray[np.floating] + alpha_collection: NDArray[np.floating] - @property - def omega_collection(self) -> NDArray: ... + internal_forces: NDArray[np.floating] + internal_torques: NDArray[np.floating] - @property - def alpha_collection(self) -> NDArray: ... + external_forces: NDArray[np.floating] + external_torques: NDArray[np.floating] - @property - def external_forces(self) -> NDArray: ... + def compute_internal_forces_and_torques(self, time: np.floating) -> None: ... - @property - def external_torques(self) -> NDArray: ... + def update_accelerations(self, time: np.floating) -> None: ... - def update_internal_forces_and_torques(self, time: np.floating) -> None: ... + def zeroed_out_external_forces_and_torques(self, time: np.floating) -> None: ... class SymplecticSystemProtocol(SystemProtocol, Protocol): @@ -55,18 +62,15 @@ class SymplecticSystemProtocol(SystemProtocol, Protocol): Protocol for system with symplectic state variables """ + v_w_collection: NDArray[np.floating] + dvdt_dwdt_collection: NDArray[np.floating] + @property def kinematic_states(self) -> _KinematicState: ... @property def dynamic_states(self) -> _DynamicState: ... - @property - def rate_collection(self) -> NDArray[np.floating]: ... - - @property - def dvdt_dwdt_collection(self) -> NDArray[np.floating]: ... - def kinematic_rates( self, time: np.floating, prefac: np.floating ) -> tuple[NDArray[np.floating], NDArray[np.floating]]: ... From bf5306aff9b6ee976179304945e7c6491408d661 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sat, 15 Jun 2024 00:03:20 +0900 Subject: [PATCH 079/134] typing: minor fix that reference RodType --- elastica/memory_block/memory_block_rod.py | 17 ++++++------- elastica/modules/base_system.py | 6 ++--- elastica/modules/memory_block.py | 9 ++++--- elastica/rigidbody/rigid_body.py | 4 ++- elastica/surface/plane.py | 6 ++--- elastica/surface/surface_base.py | 5 ++-- elastica/typing.py | 31 +++++------------------ elastica/utils.py | 4 +-- 8 files changed, 33 insertions(+), 49 deletions(-) diff --git a/elastica/memory_block/memory_block_rod.py b/elastica/memory_block/memory_block_rod.py index fcbd961de..4a1e8bd9f 100644 --- a/elastica/memory_block/memory_block_rod.py +++ b/elastica/memory_block/memory_block_rod.py @@ -1,7 +1,7 @@ __doc__ = """Create block-structure class for collection of Cosserat rod systems.""" import numpy as np from typing import Literal, Callable -from elastica.typing import SystemIdxType +from elastica.typing import SystemIdxType, RodType from elastica.memory_block.memory_block_rod_base import ( make_block_memory_metadata, make_block_memory_periodic_boundary_metadata, @@ -12,7 +12,6 @@ CosseratRod, _compute_sigma_kappa_for_blockstructure, ) -from elastica.rod.rod_base import RodBase from elastica._synchronize_periodic_boundary import ( _synchronize_periodic_boundary_of_vector_collection, _synchronize_periodic_boundary_of_scalar_collection, @@ -31,7 +30,7 @@ class MemoryBlockCosseratRod(CosseratRod, _RodSymplecticStepperMixin): """ def __init__( - self, systems: list[RodBase], system_idx_list: list[SystemIdxType] + self, systems: list[RodType], system_idx_list: list[SystemIdxType] ) -> None: # separate straight and ring rods @@ -201,7 +200,7 @@ def __init__( # Initialize the mixin class for symplectic time-stepper. _RodSymplecticStepperMixin.__init__(self) - def _allocate_block_variables_in_nodes(self, systems: list[RodBase]) -> None: + def _allocate_block_variables_in_nodes(self, systems: list[RodType]) -> None: """ This function takes system collection and allocates the variables on node for block-structure and references allocated variables back to the @@ -251,7 +250,7 @@ def _allocate_block_variables_in_nodes(self, systems: list[RodBase]) -> None: value_type="vector", ) - def _allocate_block_variables_in_elements(self, systems: list[RodBase]) -> None: + def _allocate_block_variables_in_elements(self, systems: list[RodType]) -> None: """ This function takes system collection and allocates the variables on elements for block-structure and references allocated variables back to the @@ -342,7 +341,7 @@ def _allocate_block_variables_in_elements(self, systems: list[RodBase]) -> None: value_type="tensor", ) - def _allocate_blocks_variables_in_voronoi(self, systems: list[RodBase]) -> None: + def _allocate_blocks_variables_in_voronoi(self, systems: list[RodType]) -> None: """ This function takes system collection and allocates the variables on voronoi for block-structure and references allocated variables back to the @@ -410,7 +409,7 @@ def _allocate_blocks_variables_in_voronoi(self, systems: list[RodBase]) -> None: ) def _allocate_blocks_variables_for_symplectic_stepper( - self, systems: list[RodBase] + self, systems: list[RodType] ) -> None: """ This function takes system collection and allocates the variables used by symplectic @@ -470,7 +469,7 @@ def _allocate_blocks_variables_for_symplectic_stepper( def _map_system_properties_to_block_memory( self, mapping_dict: dict, - systems: list[RodBase], + systems: list[RodType], block_memory: np.ndarray, domain_type: Literal["node", "element", "voronoi"], value_type: Literal["scalar", "vector", "tensor"], @@ -485,7 +484,7 @@ def _map_system_properties_to_block_memory( ---------- mapping_dict: dict Dictionary with attribute names as keys and block row index as values. - systems: list[RodBase] + systems: list[RodType] A sequence containing Cosserat rod objects to map from. block_memory: ndarray Memory block that, at the end of the method execution, contains all designated diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index 4b816b1d6..fd999665f 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -20,9 +20,9 @@ from collections.abc import MutableSequence -from elastica.rod import RodBase -from elastica.rigidbody import RigidBodyBase -from elastica.surface import SurfaceBase +from elastica.rod.rod_base import RodBase +from elastica.rigidbody.rigid_body import RigidBodyBase +from elastica.surface.surface_base import SurfaceBase from .memory_block import construct_memory_block_structures from .operator_group import OperatorGroupFIFO diff --git a/elastica/modules/memory_block.py b/elastica/modules/memory_block.py index a12b70d71..f13760d88 100644 --- a/elastica/modules/memory_block.py +++ b/elastica/modules/memory_block.py @@ -5,10 +5,11 @@ from elastica.typing import SystemType, SystemIdxType -from elastica.rod import RodBase -from elastica.rigidbody import RigidBodyBase -from elastica.surface import SurfaceBase -from elastica.memory_block import MemoryBlockCosseratRod, MemoryBlockRigidBody +from elastica.rod.rod_base import RodBase +from elastica.rigidbody.rigid_body import RigidBodyBase +from elastica.surface.surface_base import SurfaceBase +from elastica.memory_block.memory_block_rod import MemoryBlockCosseratRod +from elastica.memory_block.memory_block_rigid_body import MemoryBlockRigidBody def construct_memory_block_structures(systems: list[SystemType]) -> list[SystemType]: diff --git a/elastica/rigidbody/rigid_body.py b/elastica/rigidbody/rigid_body.py index 7ef340d64..5c5b76900 100644 --- a/elastica/rigidbody/rigid_body.py +++ b/elastica/rigidbody/rigid_body.py @@ -1,5 +1,7 @@ __doc__ = """""" +from typing import Type + import numpy as np from abc import ABC from elastica._linalg import _batch_matvec, _batch_cross @@ -15,7 +17,7 @@ class RigidBodyBase(ABC): """ - REQUISITE_MODULES = [] + REQUISITE_MODULES: list[Type] = [] def __init__(self): diff --git a/elastica/surface/plane.py b/elastica/surface/plane.py index a68a5be6e..82edd1a64 100644 --- a/elastica/surface/plane.py +++ b/elastica/surface/plane.py @@ -2,7 +2,6 @@ from elastica.surface.surface_base import SurfaceBase import numpy as np -from numpy.testing import assert_allclose from elastica.utils import Tolerance @@ -21,11 +20,10 @@ def __init__(self, plane_origin: np.ndarray, plane_normal: np.ndarray): Expect (3,1)-shaped array. """ - assert_allclose( + assert np.allclose( np.linalg.norm(plane_normal), 1, atol=Tolerance.atol(), - err_msg="plane normal is not a unit vector", - ) + ), "plane normal is not a unit vector" self.normal = np.asarray(plane_normal).reshape(3) self.origin = np.asarray(plane_origin).reshape(3, 1) diff --git a/elastica/surface/surface_base.py b/elastica/surface/surface_base.py index 315496120..421cb1304 100644 --- a/elastica/surface/surface_base.py +++ b/elastica/surface/surface_base.py @@ -1,4 +1,5 @@ __doc__ = """Base class for surfaces""" +from typing import Type class SurfaceBase: @@ -11,9 +12,9 @@ class SurfaceBase: """ - REQUISITE_MODULES = [] + REQUISITE_MODULES: list[Type] = [] - def __init__(self): + def __init__(self) -> None: """ SurfaceBase does not take any arguments. """ diff --git a/elastica/typing.py b/elastica/typing.py index b0f1f57af..43ccb3df8 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -13,7 +13,7 @@ if TYPE_CHECKING: # Used for type hinting without circular imports # NEVER BACK-IMPORT ANY ELASTICA MODULES HERE - from .rod import RodBase + from .rod.protocol import CosseratRodProtocol from .rigidbody import RigidBodyBase from .surface import SurfaceBase from .modules.base_system import BaseSystemCollection @@ -26,6 +26,9 @@ SymplecticStepperProtocol, MemoryProtocol, ) + from memory_block.protocol import ( + BlockCosseratRodProtocol, + ) # , BlockRigidBodyProtocol # Modules Base Classes from .boundary_conditions import FreeBC @@ -34,28 +37,6 @@ from .dissipation import DamperBase from .external_forces import NoForces from .joint import FreeJoint -# else: -# RodBase = None -# RigidBodyBase = None -# SurfaceBase = None -# -# SystemCollectionProtocol = None -# -# State = "State" -# SymplecticSystemProtocol = None -# ExplicitSystemProtocol = None -# -# StepperProtocol = None -# SymplecticStepperProtocol = None -# MemoryProtocol = None -# -# # Modules Base Classes -# FreeBC = None -# CallBackBaseClass = None -# NoContact = None -# DamperBase = None -# NoForces = None -# FreeJoint = None SystemType: TypeAlias = "SymplecticSystemProtocol" # | ExplicitSystemProtocol @@ -73,11 +54,13 @@ OperatorType: TypeAlias = Callable[..., Any] SteppersOperatorsType: TypeAlias = tuple[tuple[OperatorType, ...], ...] -RodType: TypeAlias = Type["RodBase"] +RodType: TypeAlias = "CosseratRodProtocol" +RigidBodyType: TypeAlias = "RigidBodyProtocol" SystemCollectionType: TypeAlias = "SystemCollectionProtocol" AllowedContactType: TypeAlias = ( SystemType | Type["SurfaceBase"] ) # FIXME: SurfaceBase needs to be treated differently +BlockType: TypeAlias = "BlockCosseratRodProtocol" # | "BlockRigidBodyProtocol" # Indexing types # TODO: Maybe just use slice?? diff --git a/elastica/utils.py b/elastica/utils.py index 81517d54d..d31cf3cdc 100644 --- a/elastica/utils.py +++ b/elastica/utils.py @@ -72,7 +72,7 @@ def value() -> Literal[3]: class Tolerance: @staticmethod - def atol() -> np.floating: + def atol() -> float: """ Static absolute tolerance method @@ -83,7 +83,7 @@ def atol() -> np.floating: return finfo(float64).eps * 1e4 @staticmethod - def rtol() -> np.floating: + def rtol() -> float: """ Static relative tolerance method From 0f053f2a1b29b1635f6a99cbf52dbaa32a4c3441 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sat, 15 Jun 2024 00:05:51 +0900 Subject: [PATCH 080/134] typing: rod module and refactor knot theory protocol --- elastica/__init__.py | 2 - elastica/rod/__init__.py | 8 -- elastica/rod/cosserat_rod.py | 137 +++++++++++++++++------------ elastica/rod/data_structures.py | 26 ++++-- elastica/rod/knot_theory.py | 33 ++----- pyproject.toml | 2 +- tests/test_rod/test_knot_theory.py | 8 +- 7 files changed, 110 insertions(+), 106 deletions(-) diff --git a/elastica/__init__.py b/elastica/__init__.py index b0ddf334c..7cdfc7f39 100644 --- a/elastica/__init__.py +++ b/elastica/__init__.py @@ -1,7 +1,5 @@ from collections import defaultdict from elastica.rod.knot_theory import ( - KnotTheory, - KnotTheoryCompatibleProtocol, compute_link, compute_twist, compute_writhe, diff --git a/elastica/rod/__init__.py b/elastica/rod/__init__.py index b9285b3cb..093ab0ebc 100644 --- a/elastica/rod/__init__.py +++ b/elastica/rod/__init__.py @@ -1,12 +1,4 @@ __doc__ = """Rod classes and its data structures """ -from elastica.rod.knot_theory import KnotTheory from elastica.rod.rod_base import RodBase -from elastica.rod.data_structures import ( - _RodSymplecticStepperMixin, - _State, - _DerivativeState, - _KinematicState, - _DynamicState, -) diff --git a/elastica/rod/cosserat_rod.py b/elastica/rod/cosserat_rod.py index 7d8a4f4af..868e9a912 100644 --- a/elastica/rod/cosserat_rod.py +++ b/elastica/rod/cosserat_rod.py @@ -1,8 +1,13 @@ __doc__ = """ Rod classes and implementation details """ +from typing import TYPE_CHECKING, Any, Optional, Type +from typing_extensions import Self +from elastica.typing import RodType +from .protocol import CosseratRodProtocol -import numpy as np from numpy.typing import NDArray + +import numpy as np import functools import numba from elastica.rod import RodBase @@ -13,18 +18,14 @@ _batch_matvec, ) from elastica._rotations import _inv_rotate -from elastica.rod.factory_function import allocate -from elastica.rod.knot_theory import KnotTheory from elastica._calculus import ( quadrature_kernel_for_block_structure, difference_kernel_for_block_structure, _difference, _average, ) -from typing import Any, Optional -from typing_extensions import Self - -from elastica.typing import RodType +from .factory_function import allocate +from .knot_theory import KnotTheory position_difference_kernel = _difference position_average = _average @@ -149,8 +150,10 @@ class CosseratRod(RodBase, KnotTheory): Rod element dilatation rates. """ + REQUISITE_MODULES: list[Type] = [] + def __init__( - self, + self: CosseratRodProtocol, n_elements: int, position: NDArray[np.floating], velocity: NDArray[np.floating], @@ -163,7 +166,7 @@ def __init__( inv_mass_second_moment_of_inertia: NDArray[np.floating], shear_matrix: NDArray[np.floating], bend_matrix: NDArray[np.floating], - density: NDArray[np.floating], + density_array: NDArray[np.floating], volume: NDArray[np.floating], mass: NDArray[np.floating], internal_forces: NDArray[np.floating], @@ -184,7 +187,8 @@ def __init__( internal_stress: NDArray[np.floating], internal_couple: NDArray[np.floating], ring_rod_flag: bool, - ): + ) -> None: + self.n_nodes = n_elements + 1 if ring_rod_flag else n_elements self.n_elems = n_elements self.position_collection = position self.velocity_collection = velocity @@ -197,7 +201,7 @@ def __init__( self.inv_mass_second_moment_of_inertia = inv_mass_second_moment_of_inertia self.shear_matrix = shear_matrix self.bend_matrix = bend_matrix - self.density = density + self.density = density_array self.volume = volume self.mass = mass self.internal_forces = internal_forces @@ -249,12 +253,12 @@ def straight_rod( start: NDArray[np.floating], direction: NDArray[np.floating], normal: NDArray[np.floating], - base_length: float, - base_radius: float, - density: float, + base_length: np.floating, + base_radius: np.floating, + density: np.floating, *, - nu: Optional[float] = None, - youngs_modulus: float, + nu: Optional[np.floating] = None, + youngs_modulus: np.floating, **kwargs: Any, ) -> Self: """ @@ -278,15 +282,15 @@ def straight_rod( Direction of the rod in 3D normal : NDArray[3, float] Normal vector of the rod in 3D - base_length : float + base_length : np.floating Total length of the rod - base_radius : float + base_radius : np.floating Uniform radius of the rod - density : float + density : np.floating Density of the rod - nu : float + nu : np.floating Damping coefficient for Rayleigh damping - youngs_modulus : float + youngs_modulus : np.floating Young's modulus **kwargs : dict, optional The "position" and/or "directors" can be overrided by passing "position" and "directors" argument. Remember, the shape of the "position" is (3,n_elements+1) and the shape of the "directors" is (3,3,n_elements). @@ -321,7 +325,7 @@ def straight_rod( inv_mass_second_moment_of_inertia, shear_matrix, bend_matrix, - density, + density_array, volume, mass, internal_forces, @@ -367,7 +371,7 @@ def straight_rod( inv_mass_second_moment_of_inertia, shear_matrix, bend_matrix, - density, + density_array, volume, mass, internal_forces, @@ -397,12 +401,12 @@ def ring_rod( ring_center_position: NDArray[np.floating], direction: NDArray[np.floating], normal: NDArray[np.floating], - base_length: float, - base_radius: float, - density: float, + base_length: np.floating, + base_radius: np.floating, + density: np.floating, *, - nu: Optional[float] = None, - youngs_modulus: float, + nu: Optional[np.floating] = None, + youngs_modulus: np.floating, **kwargs: Any, ) -> Self: """ @@ -425,15 +429,15 @@ def ring_rod( Direction of the rod in 3D normal : NDArray[3, float] Normal vector of the rod in 3D - base_length : float + base_length : np.floating Total length of the rod - base_radius : float + base_radius : np.floating Uniform radius of the rod - density : float + density : np.floating Density of the rod - nu : float + nu : np.floating Damping coefficient for Rayleigh damping - youngs_modulus : float + youngs_modulus : np.floating Young's modulus **kwargs : dict, optional The "position" and/or "directors" can be overrided by passing "position" and "directors" argument. Remember, the shape of the "position" is (3,n_elements+1) and the shape of the "directors" is (3,3,n_elements). @@ -537,10 +541,12 @@ def ring_rod( internal_couple, ring_rod_flag, ) - rod.REQUISITE_MODULE.append(Constraints) + rod.REQUISITE_MODULES.append(Constraints) return rod - def compute_internal_forces_and_torques(self, time: float) -> None: + def compute_internal_forces_and_torques( + self: CosseratRodProtocol, time: np.floating + ) -> None: """ Compute internal forces and torques. We need to compute internal forces and torques before the acceleration because they are used in interaction. Thus in order to speed up simulation, we will compute internal forces and torques @@ -549,7 +555,7 @@ def compute_internal_forces_and_torques(self, time: float) -> None: Parameters ---------- - time: float + time: np.floating current time """ @@ -595,13 +601,13 @@ def compute_internal_forces_and_torques(self, time: float) -> None: ) # Interface to time-stepper mixins (Symplectic, Explicit), which calls this method - def update_accelerations(self, time: float) -> None: + def update_accelerations(self: CosseratRodProtocol, time: np.floating) -> None: """ Updates the acceleration variables Parameters ---------- - time: float + time: np.floating current time """ @@ -617,12 +623,14 @@ def update_accelerations(self, time: float) -> None: self.dilatation, ) - def zeroed_out_external_forces_and_torques(self, time: np.floating) -> None: + def zeroed_out_external_forces_and_torques( + self: CosseratRodProtocol, time: np.floating + ) -> None: _zeroed_out_external_forces_and_torques( self.external_forces, self.external_torques ) - def compute_translational_energy(self) -> NDArray[np.floating]: + def compute_translational_energy(self: CosseratRodProtocol) -> NDArray[np.floating]: """ Compute total translational energy of the rod at the instance. """ @@ -636,7 +644,7 @@ def compute_translational_energy(self) -> NDArray[np.floating]: ).sum() ) - def compute_rotational_energy(self) -> NDArray[np.floating]: + def compute_rotational_energy(self: CosseratRodProtocol) -> NDArray[np.floating]: """ Compute total rotational energy of the rod at the instance. """ @@ -646,7 +654,9 @@ def compute_rotational_energy(self) -> NDArray[np.floating]: ) return 0.5 * np.einsum("ik,ik->k", self.omega_collection, J_omega_upon_e).sum() - def compute_velocity_center_of_mass(self) -> NDArray[np.floating]: + def compute_velocity_center_of_mass( + self: CosseratRodProtocol, + ) -> NDArray[np.floating]: """ Compute velocity center of mass of the rod at the instance. """ @@ -655,7 +665,9 @@ def compute_velocity_center_of_mass(self) -> NDArray[np.floating]: return sum_mass_times_velocity / self.mass.sum() - def compute_position_center_of_mass(self) -> NDArray[np.floating]: + def compute_position_center_of_mass( + self: CosseratRodProtocol, + ) -> NDArray[np.floating]: """ Compute position center of mass of the rod at the instance. """ @@ -664,7 +676,7 @@ def compute_position_center_of_mass(self) -> NDArray[np.floating]: return sum_mass_times_position / self.mass.sum() - def compute_bending_energy(self) -> NDArray[np.floating]: + def compute_bending_energy(self: CosseratRodProtocol) -> NDArray[np.floating]: """ Compute total bending energy of the rod at the instance. """ @@ -680,7 +692,7 @@ def compute_bending_energy(self) -> NDArray[np.floating]: ).sum() ) - def compute_shear_energy(self) -> NDArray[np.floating]: + def compute_shear_energy(self: CosseratRodProtocol) -> NDArray[np.floating]: """ Compute total shear energy of the rod at the instance. """ @@ -697,7 +709,7 @@ def compute_shear_energy(self) -> NDArray[np.floating]: # Below is the numba-implementation of Cosserat Rod equations. They don't need to be visible by users. -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _compute_geometry_from_state( position_collection: NDArray[np.floating], volume: NDArray[np.floating], @@ -725,7 +737,7 @@ def _compute_geometry_from_state( radius[k] = np.sqrt(volume[k] / lengths[k] / np.pi) -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _compute_all_dilatations( position_collection: NDArray[np.floating], volume: NDArray[np.floating], @@ -755,7 +767,7 @@ def _compute_all_dilatations( voronoi_dilatation[k] = voronoi_lengths[k] / rest_voronoi_lengths[k] -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _compute_dilatation_rate( position_collection: NDArray[np.floating], velocity_collection: NDArray[np.floating], @@ -786,7 +798,7 @@ def _compute_dilatation_rate( ) -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _compute_shear_stretch_strains( position_collection: NDArray[np.floating], volume: NDArray[np.floating], @@ -821,7 +833,7 @@ def _compute_shear_stretch_strains( sigma[:] = dilatation * _batch_matvec(director_collection, tangents) - z_vector -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _compute_internal_shear_stretch_stresses_from_model( position_collection: NDArray[np.floating], volume: NDArray[np.floating], @@ -861,7 +873,7 @@ def _compute_internal_shear_stretch_stresses_from_model( internal_stress[:] = _batch_matvec(shear_matrix, sigma - rest_sigma) -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _compute_bending_twist_strains( director_collection: NDArray[np.floating], rest_voronoi_lengths: NDArray[np.floating], @@ -878,7 +890,7 @@ def _compute_bending_twist_strains( kappa[2, k] = temp[2, k] / rest_voronoi_lengths[k] -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _compute_internal_bending_twist_stresses_from_model( director_collection: NDArray[np.floating], rest_voronoi_lengths: NDArray[np.floating], @@ -907,7 +919,7 @@ def _compute_internal_bending_twist_stresses_from_model( internal_couple[:] = _batch_matvec(bend_matrix, temp) -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _compute_internal_forces( position_collection: NDArray[np.floating], volume: NDArray[np.floating], @@ -968,7 +980,7 @@ def _compute_internal_forces( ) -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _compute_internal_torques( position_collection: NDArray[np.floating], velocity_collection: NDArray[np.floating], @@ -1057,7 +1069,7 @@ def _compute_internal_torques( ) -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _update_accelerations( acceleration_collection: NDArray[np.floating], internal_forces: NDArray[np.floating], @@ -1092,7 +1104,7 @@ def _update_accelerations( ) * dilatation[k] -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _zeroed_out_external_forces_and_torques( external_forces: NDArray[np.floating], external_torques: NDArray[np.floating] ) -> None: @@ -1115,3 +1127,16 @@ def _zeroed_out_external_forces_and_torques( for i in range(3): for k in range(n_elems): external_torques[i, k] = 0.0 + + +if TYPE_CHECKING: + _: CosseratRodProtocol = CosseratRod.straight_rod( + 3, + np.zeros(3), + np.array([0, 1, 0]), + np.array([0, 0, 1]), + np.float64(1.0), + np.float64(0.1), + np.float64(1.0), + youngs_modulus=np.float64(1.0), + ) diff --git a/elastica/rod/data_structures.py b/elastica/rod/data_structures.py index 35e2520a9..dcbbee20c 100644 --- a/elastica/rod/data_structures.py +++ b/elastica/rod/data_structures.py @@ -1,6 +1,6 @@ __doc__ = "Data structure wrapper for rod components" -from typing import Optional +from typing import TYPE_CHECKING, Optional from typing_extensions import Self import numpy as np from numpy.typing import NDArray @@ -8,6 +8,10 @@ from elastica._rotations import _get_rotation_matrix, _rotate from elastica._linalg import _batch_matmul +if TYPE_CHECKING: + from elastica.memory_block.protocol import BlockCosseratRodProtocol +else: + BlockCosseratRodProtocol = "BlockCosseratRodProtocol" # FIXME : Explicit Stepper doesn't work as States lose the # views they initially had when working with a timestepper. @@ -46,7 +50,7 @@ class _RodSymplecticStepperMixin: - def __init__(self) -> None: + def __init__(self: BlockCosseratRodProtocol) -> None: self.kinematic_states = _KinematicState( self.position_collection, self.director_collection ) @@ -63,18 +67,22 @@ def __init__(self) -> None: # is another function self.kinematic_rates = self.dynamic_states.kinematic_rates - def update_internal_forces_and_torques(self, time: np.floating) -> None: + def update_internal_forces_and_torques( + self: BlockCosseratRodProtocol, time: np.floating + ) -> None: self.compute_internal_forces_and_torques(time) def dynamic_rates( - self, + self: BlockCosseratRodProtocol, time: np.floating, prefac: np.floating, ) -> NDArray[np.floating]: self.update_accelerations(time) return self.dynamic_states.dynamic_rates(time, prefac) - def reset_external_forces_and_torques(self, time: np.floating) -> None: + def reset_external_forces_and_torques( + self: BlockCosseratRodProtocol, time: np.floating + ) -> None: self.zeroed_out_external_forces_and_torques(time) @@ -127,7 +135,7 @@ def _bootstrap_from_data( position = np.ndarray.view(vector_states[..., :n_nodes]) directors = np.ndarray.view(matrix_states) v_w_dvdt_dwdt = np.ndarray.view(vector_states[..., n_nodes:]) - output = () + output: tuple = () if stepper_type == "explicit": v_w_states = np.ndarray.view(vector_states[..., n_nodes : 3 * n_nodes - 1]) output += ( @@ -332,7 +340,7 @@ def __init__( super(_DerivativeState, self).__init__() self.rate_collection = rate_collection_view - def __rmul__(self, scalar: np.floating) -> NDArray[np.floating]: + def __rmul__(self, scalar: np.floating) -> NDArray[np.floating]: # type: ignore """overloaded scalar * self, Parameters @@ -517,7 +525,9 @@ def kinematic_rates( # Comes from kin_state -> (x,Q) += dt * (v,w) <- First part of dyn_state return self.velocity_collection, self.omega_collection - def dynamic_rates(self, time: np.floating, prefac: np.floating): + def dynamic_rates( + self, time: np.floating, prefac: np.floating + ) -> NDArray[np.floating]: """Yields dynamic rates to add to with _DynamicState Returns ------- diff --git a/elastica/rod/knot_theory.py b/elastica/rod/knot_theory.py index 4f959ce06..404131e35 100644 --- a/elastica/rod/knot_theory.py +++ b/elastica/rod/knot_theory.py @@ -10,7 +10,7 @@ The details discussion is included in `N Charles et. al. PRL (2019) `_. """ -from typing import Protocol, Union +pass from numba import njit import numpy as np @@ -19,24 +19,7 @@ from elastica.rod.rod_base import RodBase from elastica._linalg import _batch_norm, _batch_dot, _batch_cross - -class KnotTheoryCompatibleProtocol(Protocol): - """KnotTheoryCompatibleProtocol - - Required properties to use KnotTheory mixin - """ - - @property - def position_collection(self) -> np.ndarray: ... - - @property - def director_collection(self) -> np.ndarray: ... - - @property - def radius(self) -> np.ndarray: ... - - @property - def base_length(self) -> np.ndarray: ... +from .protocol import CosseratRodProtocol class KnotTheory: @@ -77,9 +60,7 @@ def __init__(self) -> None: """ - MIXIN_PROTOCOL = Union[RodBase, KnotTheoryCompatibleProtocol] - - def compute_twist(self: MIXIN_PROTOCOL) -> NDArray[np.floating]: + def compute_twist(self: CosseratRodProtocol) -> NDArray[np.floating]: """ See :ref:`api/rods:Knot Theory (Mixin)` for the detail. """ @@ -90,9 +71,9 @@ def compute_twist(self: MIXIN_PROTOCOL) -> NDArray[np.floating]: return total_twist[0] def compute_writhe( - self: MIXIN_PROTOCOL, + self: CosseratRodProtocol, type_of_additional_segment: str = "next_tangent", - alpha: np.floating = 1.0, + alpha: float = 1.0, ) -> NDArray[np.floating]: """ See :ref:`api/rods:Knot Theory (Mixin)` for the detail. @@ -113,9 +94,9 @@ def compute_writhe( )[0] def compute_link( - self: MIXIN_PROTOCOL, + self: CosseratRodProtocol, type_of_additional_segment: str = "next_tangent", - alpha: np.floating = 1.0, + alpha: float = 1.0, ) -> NDArray[np.floating]: """ See :ref:`api/rods:Knot Theory (Mixin)` for the detail. diff --git a/pyproject.toml b/pyproject.toml index f665dd07e..71b03f98d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -137,7 +137,7 @@ strict_equality = true strict_optional = true warn_no_return = true warn_redundant_casts = true -warn_return_any = true +warn_return_any = false warn_unreachable = false warn_unused_configs = true warn_unused_ignores = false diff --git a/tests/test_rod/test_knot_theory.py b/tests/test_rod/test_knot_theory.py index a5630d73a..a74cc252b 100644 --- a/tests/test_rod/test_knot_theory.py +++ b/tests/test_rod/test_knot_theory.py @@ -10,13 +10,14 @@ from elastica.rod.rod_base import RodBase from elastica.rod.knot_theory import ( - KnotTheoryCompatibleProtocol, compute_twist, compute_writhe, compute_link, _compute_additional_segment, ) +from elastica.rod.protocol import CosseratRodProtocol + @pytest.fixture def knot_theory(): @@ -28,7 +29,7 @@ def knot_theory(): def test_knot_theory_protocol(): # To clear the protocol test coverage with pytest.raises(TypeError) as e_info: - protocol = KnotTheoryCompatibleProtocol() + protocol = CosseratRodProtocol() assert "cannot be instantiated" in e_info @@ -39,9 +40,6 @@ def __init__(self): self.radius = np.random.randn(MaxDimension.value(), self.n_elems) rod = TestRodWithKnotTheory() - assert hasattr( - rod, "MIXIN_PROTOCOL" - ), "Expected to mix-in variables: MIXIN_PROTOCOL" assert hasattr( rod, "compute_writhe" ), "Expected to mix-in functionals into the rod class: compute_writhe" From d7da64525e2ceacd6c3c246407bbf4ec42e708e8 Mon Sep 17 00:00:00 2001 From: Songyuan Cui Date: Fri, 14 Jun 2024 23:31:31 +0800 Subject: [PATCH 081/134] refactor: remove unused bootstrap code from rigid body data_structure --- elastica/rigidbody/data_structures.py | 458 -------------------------- 1 file changed, 458 deletions(-) diff --git a/elastica/rigidbody/data_structures.py b/elastica/rigidbody/data_structures.py index 6cf64a6d6..d6edb4f22 100644 --- a/elastica/rigidbody/data_structures.py +++ b/elastica/rigidbody/data_structures.py @@ -51,461 +51,3 @@ def __init__(self) -> None: def update_internal_forces_and_torques(self, *args: Any, **kwargs: Any) -> None: pass - - -# def _bootstrap_from_data(stepper_type: str, n_elems: int, vector_states, matrix_states): -# """Returns states wrapping numpy arrays based on the time-stepping algorithm -# -# Convenience method that takes in rod internal (raw np.ndarray) data, create views -# (references) from it, and outputs State classes that are used in the time-stepping -# algorithm. This means that modifying the state modifies the internal data! -# -# Parameters -# ---------- -# stepper_type : str (likely to change in future), representing stepper type -# Allowed parameters are ['explicit', 'symplectic'] -# n_elems : int, number of rod elements -# vector_states : np.ndarray of shape (dim, *) with the following structure -# `vector_states` = [`position`,`velocity`,`omega`,`acceleration`,`angular acceleration`] -# `n_nodes = n_elems = 1` -# `position = 0 -> n_nodes , size = n_nodes` -# `velocity = n_nodes -> 2 * n_nodes, size = n_nodes` -# `omega = 2 * n_nodes -> 2 * n_nodes + nelem, size = nelem` -# `acceleration = 2 * n_nodes + nelem -> 3 * n_nodes + nelem, size = n_nodes` -# `angular acceleration = 3 * n_nodes + nelem -> 3 * n_nodes + 2 * nelem, size = n_elems` -# matrix_states : np.ndarray of shape (dim, dim, n_elems) containing the directors -# -# Returns -# ------- -# output : tuple of len 8 containing -# (state, derivative_state, position, directors, velocity, omega, acceleration, alpha) -# derivative_state carries rate information -# -# """ -# n_nodes = n_elems -# position = np.ndarray.view(vector_states[..., :n_nodes]) -# directors = np.ndarray.view(matrix_states) -# v_w_dvdt_dwdt = np.ndarray.view(vector_states[..., n_nodes:]) -# output = () -# TODO: -# 11/01/2020: Extend Rigid body data structures for Explicit steppers -# 12/20/2021: Future work! If explicit stepper is gonna be dropped, remove. -# if stepper_type == "explicit": -# v_w_states = np.ndarray.view(vector_states[..., n_nodes : n_nodes + 1]) -# output += ( -# _State(n_elems, position, directors, v_w_states), -# _DerivativeState(n_elems, v_w_dvdt_dwdt), -# ) -# elif stepper_type == "symplectic": -# output += ( -# _KinematicState(n_elems, position, directors), -# _DynamicState(n_elems, v_w_dvdt_dwdt), -# ) -# else: -# return -# output += ( -# _KinematicState(n_elems, position, directors), -# _DynamicState(n_elems, v_w_dvdt_dwdt), -# ) -# -# n_velocity_end = n_nodes + n_nodes -# velocity = np.ndarray.view(vector_states[..., n_nodes:n_velocity_end]) -# -# n_omega_end = n_velocity_end + n_elems -# omega = np.ndarray.view(vector_states[..., n_velocity_end:n_omega_end]) -# -# n_acceleration_end = n_omega_end + n_nodes -# acceleration = np.ndarray.view(vector_states[..., n_omega_end:n_acceleration_end]) -# -# n_alpha_end = n_acceleration_end + n_elems -# alpha = np.ndarray.view(vector_states[..., n_acceleration_end:n_alpha_end]) -# -# return output + (position, directors, velocity, omega, acceleration, alpha) - - -# """ -# Explicit stepper interface -# """ - -# TODO -# 12/20/2021: If explicit stepper is gonna be dropped, remove. -# class _State: -# """State for explicit steppers. -# -# Wraps data as state, with overloaded methods for explicit steppers -# (steppers that integrate all states in one-step/stage). -# Allows for separating implementation of stepper from actual -# addition/multiplication/other formulae used. -# """ -# -# # TODO : args, kwargs instead of hardcoding types -# def __init__( -# self, -# n_elems: int, -# position_collection_view, -# director_collection_view, -# kinematic_rate_collection_view, -# ): -# """ -# Parameters -# ---------- -# n_elems : int, number of rod elements -# position_collection_view : view of positions (or) x -# director_collection_view : view of directors (or) Q -# kinematic_rate_collection_view : view of velocity and omega (or) (v,ω) -# """ -# super(_State, self).__init__() -# self.n_nodes = n_elems -# self.n_kinematic_rates = self.n_nodes + n_elems # start of (v,ω) in (x,Q,v,ω) -# self.position_collection = position_collection_view -# self.director_collection = director_collection_view -# self.kinematic_rate_collection = kinematic_rate_collection_view -# -# def __iadd__(self, scaled_deriv_array): -# """overloaded += operator -# -# The add for directors is customized to reflect Rodrigues' rotation -# formula. -# -# Parameters -# ---------- -# scaled_deriv_array : np.ndarray containing dt * (v, ω, dv/dt, dω/dt) -# ,as returned from _DerivativeState's __mul__ method -# -# Returns -# ------- -# self : _State with inplace modified data -# -# """ -# # x += v*dt -# self.position_collection += scaled_deriv_array[..., : self.n_nodes] -# # TODO : Verify the math in this note -# r""" -# Developer Note -# -------------- -# Here the overloaded `+=` operator is exploited to perform -# matrix multiplication for the directors, which is counter- -# intutive at first. While this provides a stable interface -# to interact the rod states with the timesteppers and the -# rest of the world, the reasons behind including it here also has -# a depper mathematical significance. -# -# Firstly, position lies in the vector space corresponding to R^{3} -# and update is done this space (with the + and * operators defined -# as usual), hence the `+=` operator (or `__iadd__`) is reflected -# as `+=` operator in the position update (line 163 above). -# -# For directors rather, which lie in a restricteed R^{3} \otimes -# R^{3} tensorial space, the space with Q^T.Q = Q.Q^T = I, the + -# operator can be thought of as an equivalent `*=` update for a -# 'exponential' multiplication with a rotation matrix (e^{At}). -# . This does not correspond to the position update. However, if -# we view this in a logarithmic space the `*=` becomse the '+=' -# operator once again! After performing this `+=` operation, we -# bring it back into its original space using the exponential -# operator. So we are still indirectly doing the '+=' -# update. -# -# To avoid all this hassle with the operators and spaces, we simply define -# '+=' or '__iadd__' in the case of directors as an equivalent -# '*=' (matrix multiply) with the RHS below. -# """ -# # TODO Q *= exp(w*dt) , whats' the formua again? -# # TODO the scale factor 1.0 does not seem to be necessary, although -# # we perform more work in the present framework (muliply dt to entire vector, then take -# # norm) rather than vector norm then multiple by dt (1/3 operation costs) -# # TODO optimize (somehow) extra copy away : if we don't make a copy -# # its even more slower, maybe due to aliasing effects -# np.einsum( -# "ijk,jlk->ilk", -# _get_rotation_matrix( -# 1.0, scaled_deriv_array[..., self.n_nodes : self.n_kinematic_rates] -# ), -# self.director_collection.copy(), -# out=self.director_collection, -# ) -# # (v,ω) += (dv/dt, dω/dt)*dt -# self.kinematic_rate_collection += scaled_deriv_array[ -# ..., self.n_kinematic_rates : -# ] -# return self -# -# def __add__(self, scaled_derivative_state): -# """overloaded + operator, useful in state.k1 = state + dt * deriv_state -# -# The add for directors is customized to reflect Rodrigues' rotation -# formula. -# -# Parameters -# ---------- -# scaled_derivative_state : np.ndarray with dt * (v, ω, dv/dt, dω/dt) -# ,as returned from _DerivativeState's __mul__ method -# -# Returns -# ------- -# state : new _State object with modified data (copied) -# -# Caveats -# ------- -# Note that the argument is not a `other` _State object but is rather -# assumed to be a `np.ndarray` from calling _DerivativeState's __mul__ -# method. This reflects the most common use-case in time-steppers -# -# """ -# # x += v*dt -# position_collection = ( -# self.position_collection + scaled_derivative_state[..., : self.n_nodes] -# ) -# # Devs : see `_State.__iadd__` for reasons why we do matmul here -# director_collection = _rotate( -# self.director_collection, -# 1.0, -# scaled_derivative_state[..., self.n_nodes : self.n_kinematic_rates], -# ) -# # (v,ω) += (dv/dt, dω/dt)*dt -# kinematic_rate_collection = ( -# self.kinematic_rate_collection -# + scaled_derivative_state[..., self.n_kinematic_rates :] -# ) -# return _State( -# self.n_nodes, -# position_collection, -# director_collection, -# kinematic_rate_collection, -# ) -# -# -# class _DerivativeState: -# """TimeDerivative of States for explicit steppers. -# -# Wraps time-derivative data as state, with overloaded methods for -# explicit steppers (steppers that integrate all states in one-step/stage). -# Allows for separating implementation of stepper from actual addition -# /multiplication used. -# """ -# -# def __init__(self, _unused_n_elems: int, rate_collection_view): -# """ -# Parameters -# ---------- -# _unused_n_elems : int, number of elements (unused, kept for -# compatibility with `_bootstrap_from_data`) -# rate_collection_view : np.ndarray containing (v, ω, dv/dt, dω/dt) -# """ -# super(_DerivativeState, self).__init__() -# self.rate_collection = rate_collection_view -# -# def __rmul__(self, scalar): -# """overloaded scalar * self, -# -# Parameters -# ---------- -# scalar : float, typically dt (the time-step) -# -# Returns -# ------- -# output : np.ndarray containing (v*dt, ω*dt, dv/dt*dt, dω/dt*dt) -# -# Caveats -# ------- -# Returns a np.ndarray and not a State object (as one expects). -# Returning a State here with (v*dt, ω*dt, dv/dt*dt, dω/dt*dt) as members -# is possible but it's less efficient, especially because this is hot -# piece of code -# """ -# """ -# Developer Note -# -------------- -# -# Q : Why do we need to overload operators here? -# -# The Derivative class naturally doesn't have a `mul` overloaded -# operator. That means if this method is not present, -# doing something like -# ``` -# ds = _DerivativeState(...) -# new_state = 2 * ds -# ``` -# will throw an error. Note that you can do something like -# ``` -# ds = _DerivativeState(...) -# new_state = 2 * ds.rate_collection -# ``` -# but this is hacky, as we are exposing the members outside, -# in the calling scope (defeats encapsulation and hiding). -# The point of having this class is that it works -# well with the time-stepper (where we only use `+` and `*` -# operations on the State/DerivativeState like above, -# i.e. `state = dt * derivative_state` and not something like -# `state = dt * derivative_state.rate_collection`). -# It also provides an interface for anything outside -# the `Rod` system as a whole. -# """ -# return scalar * self.rate_collection -# -# def __mul__(self, scalar): -# """overloaded self * scalar -# -# TODO Check if this pattern (forwarding to __mul__) has -# any disdvantages apart from extra function call penalty -# -# Parameters -# ---------- -# scalar : float, typically dt (the time-step) -# -# Returns -# ------- -# output : np.ndarray containing (v*dt, ω*dt, dv/dt*dt, dω/dt*dt) -# -# """ -# return self.__rmul__(scalar) - -# """ -# Symplectic stepper interface -# """ - - -# TODO: Maybe considerg removing. We no longer use bootstrap to initialize. -# RigidBodySymplecticStepperMixin is now derived from RodSymplecticStepperMixin -# class _KinematicState: -# """State storing (x,Q) for symplectic steppers. -# -# Wraps data as state, with overloaded methods for symplectic steppers. -# Allows for separating implementation of stepper from actual -# addition/multiplication/other formulae used. -# -# Symplectic steppers rely only on in-place modifications to state and so -# only these methods are provided. -# """ -# -# def __init__(self, position_collection_view, director_collection_view): -# """ -# Parameters -# ---------- -# n_elems : int, number of rod elements -# position_collection_view : view of positions (or) x -# director_collection_view : view of directors (or) Q -# """ -# self.position_collection = position_collection_view -# self.director_collection = director_collection_view -# -# def __iadd__(self, scaled_deriv_array): -# """overloaded += operator -# -# The add for directors is customized to reflect Rodrigues' rotation -# formula. -# -# Parameters -# ---------- -# scaled_deriv_array : np.ndarray containing dt * (v, ω), -# as retured from _DynamicState's `kinematic_rates` method -# -# Returns -# ------- -# self : _KinematicState instance with inplace modified data -# -# Caveats -# ------- -# Takes a np.ndarray and not a _KinematicState object (as one expects). -# This is done for efficiency reasons, see _DynamicState's `kinematic_rates` -# method -# """ -# velocity_collection = scaled_deriv_array[0] -# omega_collection = scaled_deriv_array[1] -# # x += v*dt -# self.position_collection += velocity_collection -# # Devs : see `_State.__iadd__` for reasons why we do matmul here -# np.einsum( -# "ijk,jlk->ilk", -# _get_rotation_matrix(1.0, omega_collection), -# self.director_collection.copy(), -# out=self.director_collection, -# ) -# return self -# -# -# class _DynamicState: -# """State storing (v,ω, dv/dt, dω/dt) for symplectic steppers. -# -# Wraps data as state, with overloaded methods for symplectic steppers. -# Allows for separating implementation of stepper from actual -# addition/multiplication/other formulae used. -# -# Symplectic steppers rely only on in-place modifications to state and so -# only these methods are provided. -# """ -# -# def __init__( -# self, -# v_w_collection, -# dvdt_dwdt_collection, -# velocity_collection, -# omega_collection, -# ): -# """ -# -# Parameters -# ---------- -# n_elems : int, number of rod elements -# rate_collection_view : np.ndarray containing (v, ω, dv/dt, dω/dt) -# """ -# super(_DynamicState, self).__init__() -# # Limit at which (v, w) end -# self.rate_collection = v_w_collection -# self.dvdt_dwdt_collection = dvdt_dwdt_collection -# self.velocity_collection = velocity_collection -# self.omega_collection = omega_collection -# -# def __iadd__(self, scaled_second_deriv_array): -# """overloaded += operator, updating dynamic_rates -# -# Parameters -# ---------- -# scaled_second_deriv_array : np.ndarray containing dt * (dvdt, dωdt), -# as retured from _DynamicState's `dynamic_rates` method -# -# Returns -# ------- -# self : _DynamicState instance with inplace modified data -# -# Caveats -# ------- -# Takes a np.ndarray and not a _DynamicState object (as one expects). -# This is done for efficiency reasons, see `dynamic_rates`. -# """ -# # Always goes in LHS : that means the update is on the rates alone -# # (v,ω) += dt * (dv/dt, dω/dt) -> self.dynamic_rates -# self.rate_collection += scaled_second_deriv_array -# return self -# -# def kinematic_rates(self, time, prefac, *args, **kwargs): -# """Yields kinematic rates to interact with _KinematicState -# -# Returns -# ------- -# v_and_omega : np.ndarray consisting of (v,ω) -# -# Caveats -# ------- -# Doesn't return a _KinematicState with (dt*v, dt*w) as members, -# as one expects the _Kinematic __add__ operator to interact -# with another _KinematicState. This is done for efficiency purposes. -# """ -# # RHS functino call, gives v,w so that -# # Comes from kin_state -> (x,Q) += dt * (v,w) <- First part of dyn_state -# return prefac * self.velocity_collection, prefac * self.omega_collection -# -# def dynamic_rates(self, time, prefac, *args, **kwargs): -# """Yields dynamic rates to add to with _DynamicState -# -# Returns -# ------- -# acc_and_alpha : np.ndarray consisting of (dv/dt,dω/dt) -# -# Caveats -# ------- -# Doesn't return a _DynamicState with (dt*v, dt*w) as members, -# as one expects the _Dynamic __add__ operator to interact -# with another _DynamicState. This is done for efficiency purposes. -# """ -# return prefac * self.dvdt_dwdt_collection From 51f01a9c410d70fd5c051faec1593cd1139057fb Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Thu, 13 Jun 2024 22:38:24 +0900 Subject: [PATCH 082/134] typing: system utility --- elastica/systems/__init__.py | 45 ++++++++++++++++++++---------------- elastica/typing.py | 2 +- 2 files changed, 26 insertions(+), 21 deletions(-) diff --git a/elastica/systems/__init__.py b/elastica/systems/__init__.py index 16a17ff20..ca2bd8306 100644 --- a/elastica/systems/__init__.py +++ b/elastica/systems/__init__.py @@ -1,3 +1,10 @@ +from typing import Iterator, TypeVar, Generic, Type +from elastica.timestepper.protocol import ExplicitStepperProtocol +from elastica.typing import SystemCollectionType + +from copy import copy + + def is_system_a_collection(system: object) -> bool: # Check if system is a "collection" of smaller systems # by checking for the [] method @@ -38,7 +45,10 @@ def is_system_a_collection(system: object) -> bool: ) -def make_memory_for_explicit_stepper(stepper, system): +# TODO: Use MemoryProtocol +def make_memory_for_explicit_stepper( + stepper: ExplicitStepperProtocol, system: SystemCollectionType +) -> "MemoryCollection": # TODO Automated logic (class creation, memory management logic) agnostic of stepper details (RK, AB etc.) from elastica.timestepper.explicit_steppers import ( @@ -46,14 +56,14 @@ def make_memory_for_explicit_stepper(stepper, system): EulerForward, ) - is_this_system_a_collection = is_system_a_collection(system) + # is_this_system_a_collection = is_system_a_collection(system) + memory_cls: Type if RungeKutta4 in stepper.__class__.mro(): # Bad way of doing it, introduces tight coupling # this should rather be taken from the class itself class MemoryRungeKutta4: - def __init__(self): - super(MemoryRungeKutta4, self).__init__() + def __init__(self) -> None: self.initial_state = None self.k_1 = None self.k_2 = None @@ -62,19 +72,17 @@ def __init__(self): memory_cls = MemoryRungeKutta4 elif EulerForward in stepper.__class__.mro(): - memory_cls = NotImplementedError + raise NotImplementedError("Euler Forward not implemented") else: - # TODO Memory allocation for other integrators raise NotImplementedError("Making memory for other types not supported") - return ( - MemoryCollection(memory_cls(), len(system)) - if is_this_system_a_collection - else memory_cls() - ) + return MemoryCollection(memory_cls(), len(system)) -class MemoryCollection: +M = TypeVar("M", bound="MemoryCollection") + + +class MemoryCollection(Generic[M]): """Slots of memories for timestepper in a cohesive unit. A `MemoryCollection` object is meant to be used in conjunction @@ -91,24 +99,21 @@ class MemoryCollection: every stage. """ - def __init__(self, memory, n_memory_slots): + def __init__(self, memory: M, n_memory_slots: int): super(MemoryCollection, self).__init__() - self.__memories = [None] * n_memory_slots - - from copy import copy - + self.__memories: list[M] = [] for i_slot in range(n_memory_slots - 1): self.__memories[i_slot] = copy(memory) # Save final copy self.__memories[-1] = memory - def __getitem__(self, idx): + def __getitem__(self, idx: int) -> M: return self.__memories[idx] - def __len__(self): + def __len__(self) -> int: return len(self.__memories) - def __iter__(self): + def __iter__(self) -> Iterator[M]: return self.__memories.__iter__() diff --git a/elastica/typing.py b/elastica/typing.py index fd1337fe5..d7d26f514 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -75,7 +75,7 @@ SteppersOperatorsType: TypeAlias = tuple[tuple[OperatorType, ...], ...] # tuple[Union[PrefactorOperatorType, StepOperatorType, NoOpType, np.floating], ...], ... # Explicit stepper -# ExplicitStageOperatorType = Callable[ +# ExplicitStageOperatorType: TypeAlias = Callable[ # [ # SymplecticStepperProtocol, # ExplicitSystemProtocol, From d9793047d6526d71b40c75efbd1605a70238893a Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Thu, 13 Jun 2024 22:43:26 +0900 Subject: [PATCH 083/134] feat: Euler Forward memory --- elastica/systems/__init__.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/elastica/systems/__init__.py b/elastica/systems/__init__.py index ca2bd8306..246629c39 100644 --- a/elastica/systems/__init__.py +++ b/elastica/systems/__init__.py @@ -45,6 +45,9 @@ def is_system_a_collection(system: object) -> bool: ) +# FIXME: Move memory related functions to separate module or as part of the timestepper + + # TODO: Use MemoryProtocol def make_memory_for_explicit_stepper( stepper: ExplicitStepperProtocol, system: SystemCollectionType @@ -72,7 +75,13 @@ def __init__(self) -> None: memory_cls = MemoryRungeKutta4 elif EulerForward in stepper.__class__.mro(): - raise NotImplementedError("Euler Forward not implemented") + + class MemoryEulerForward: + def __init__(self) -> None: + self.initial_state = None + self.k = None + + memory_cls = MemoryEulerForward else: raise NotImplementedError("Making memory for other types not supported") @@ -103,11 +112,9 @@ def __init__(self, memory: M, n_memory_slots: int): super(MemoryCollection, self).__init__() self.__memories: list[M] = [] - for i_slot in range(n_memory_slots - 1): - self.__memories[i_slot] = copy(memory) - - # Save final copy - self.__memories[-1] = memory + for _ in range(n_memory_slots - 1): + self.__memories.append(copy(memory)) + self.__memories.append(memory) def __getitem__(self, idx: int) -> M: return self.__memories[idx] From 3e76a3826721ea706fc89de62a7725fcd0ac2faa Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 16 Jun 2024 00:48:48 +0900 Subject: [PATCH 084/134] typing: contact related modules --- elastica/contact_forces.py | 389 ++++++------------ elastica/contact_utils.py | 39 +- elastica/interaction.py | 6 +- .../snake_contact.py | 34 +- tests/test_contact_classes.py | 201 ++++----- tests/test_modules/test_contact.py | 39 +- 6 files changed, 271 insertions(+), 437 deletions(-) diff --git a/elastica/contact_forces.py b/elastica/contact_forces.py index 9e9f5a2de..22adcb3e9 100644 --- a/elastica/contact_forces.py +++ b/elastica/contact_forces.py @@ -1,10 +1,13 @@ __doc__ = """ Numba implementation module containing contact between rods and rigid bodies and other rods rigid bodies or surfaces.""" -from typing import Optional -from elastica.typing import RodType, SystemType, AllowedContactType -from elastica.rod import RodBase -from elastica.rigidbody import Cylinder, Sphere -from elastica.surface import Plane +from typing import TypeVar, Generic, Type +from elastica.typing import RodType, SystemType, SurfaceType + +from elastica.rod.rod_base import RodBase +from elastica.rigidbody.cylinder import Cylinder +from elastica.rigidbody.sphere import Sphere +from elastica.surface.plane import Plane +from elastica.surface.surface_base import SurfaceBase from elastica.contact_utils import ( _prune_using_aabbs_rod_cylinder, _prune_using_aabbs_rod_rod, @@ -23,7 +26,11 @@ from numpy.typing import NDArray -class NoContact: +S1 = TypeVar("S1") # TODO: Find bound +S2 = TypeVar("S2") + + +class NoContact(Generic[S1, S2]): """ This is the base class for contact applied between rod-like objects and allowed contact objects. @@ -38,51 +45,48 @@ def __init__(self) -> None: """ NoContact class does not need any input parameters. """ + pass + + @property + def _allowed_system_one(self) -> list[Type]: + # Modify this list to include the allowed system types for contact + return [RodBase] + + @property + def _allowed_system_two(self) -> list[Type]: + # Modify this list to include the allowed system types for contact + return [RodBase] def _check_systems_validity( self, - system_one: SystemType, - system_two: AllowedContactType, + system_one: S1, + system_two: S2, ) -> None: """ - This checks the contact order between a SystemType object and an AllowedContactType object, the order should follow: Rod, Rigid body, Surface. - In NoContact class, this just checks if system_two is a rod then system_one must be a rod. + Here, we check the allowed system types for contact. + For derived classes, this method can be overridden to enforce specific system types + for contact model. + """ + common_check_systems_validity(system_one, self._allowed_system_one) + common_check_systems_validity(system_two, self._allowed_system_two) - Parameters - ---------- - system_one - SystemType - system_two - AllowedContactType - """ - if issubclass(system_two.__class__, RodBase): - if not issubclass(system_one.__class__, RodBase): - raise TypeError( - "Systems provided to the contact class have incorrect order. \n" - " First system is {0} and second system is {1}. \n" - " If the first system is a rod, the second system can be a rod, rigid body or surface. \n" - " If the first system is a rigid body, the second system can be a rigid body or surface.".format( - system_one.__class__, system_two.__class__ - ) - ) + common_check_systems_identity(system_one, system_two) def apply_contact( self, - system_one: SystemType, - system_two: AllowedContactType, - ) -> Optional[tuple[NDArray[np.floating], NDArray[np.intp]]]: + system_one: S1, + system_two: S2, + ) -> None: """ - Apply contact forces and torques between SystemType object and AllowedContactType object. + Apply contact forces and torques between two system object.. In NoContact class, this routine simply passes. Parameters ---------- - system_one : SystemType - Rod or rigid-body object - system_two : AllowedContactType - Rod, rigid-body, or surface object + system_one + system_two """ pass @@ -116,49 +120,14 @@ def __init__(self, k: np.floating, nu: np.floating) -> None: self.k = k self.nu = nu - def _check_systems_validity( - self, - system_one: SystemType, - system_two: AllowedContactType, - ) -> None: - """ - This checks the contact order and type of a SystemType object and an AllowedContactType object. - For the RodRodContact class both systems must be distinct rods. - - Parameters - ---------- - system_one - SystemType - system_two - AllowedContactType - """ - if not issubclass(system_one.__class__, RodBase) or not issubclass( - system_two.__class__, RodBase - ): - raise TypeError( - "Systems provided to the contact class have incorrect order. \n" - " First system is {0} and second system is {1}. \n" - " Both systems must be distinct rods".format( - system_one.__class__, system_two.__class__ - ) - ) - if system_one == system_two: - raise TypeError( - "First rod is identical to second rod. \n" - "Rods must be distinct for RodRodConact. \n" - "If you want self contact, use RodSelfContact instead" - ) - def apply_contact(self, system_one: RodType, system_two: RodType) -> None: """ Apply contact forces and torques between RodType object and RodType object. Parameters ---------- - system_one: object - Rod object. - system_two: object - Rod object. + system_one: RodType + system_two: RodType """ # First, check for a global AABB bounding box, and see whether that @@ -229,8 +198,8 @@ def __init__( self, k: np.floating, nu: np.floating, - velocity_damping_coefficient: np.floating = 0.0, - friction_coefficient: np.floating = 0.0, + velocity_damping_coefficient: float = 0.0, + friction_coefficient: float = 0.0, ) -> None: """ @@ -252,34 +221,12 @@ def __init__( self.velocity_damping_coefficient = velocity_damping_coefficient self.friction_coefficient = friction_coefficient - def _check_systems_validity( - self, - system_one: SystemType, - system_two: AllowedContactType, - ) -> None: - """ - This checks the contact order and type of a SystemType object and an AllowedContactType object. - For the RodCylinderContact class first_system should be a rod and second_system should be a cylinder. + @property + def _allowed_system_two(self) -> list[Type]: + # Modify this list to include the allowed system types for contact + return [Cylinder] - Parameters - ---------- - system_one - SystemType - system_two - AllowedContactType - """ - if not issubclass(system_one.__class__, RodBase) or not issubclass( - system_two.__class__, Cylinder - ): - raise TypeError( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a cylinder".format( - system_one.__class__, system_two.__class__ - ) - ) - - def apply_contact(self, system_one: RodType, system_two: SystemType) -> None: + def apply_contact(self, system_one: RodType, system_two: Cylinder) -> None: # First, check for a global AABB bounding box, and see whether that # intersects if _prune_using_aabbs_rod_cylinder( @@ -288,8 +235,8 @@ def apply_contact(self, system_one: RodType, system_two: SystemType) -> None: system_one.lengths, system_two.position_collection, system_two.director_collection, - system_two.radius[0], - system_two.length[0], + system_two.radius, + system_two.length, ): return @@ -356,33 +303,15 @@ def __init__(self, k: np.floating, nu: np.floating) -> None: def _check_systems_validity( self, - system_one: SystemType, - system_two: AllowedContactType, + system_one: RodType, + system_two: RodType, ) -> None: """ - This checks the contact order and type of a SystemType object and an AllowedContactType object. - For the RodSelfContact class first_system and second_system should be the same rod. - - Parameters - ---------- - system_one - SystemType - system_two - AllowedContactType + Overriding the base class method to check if the two systems are identical. """ - if ( - not issubclass(system_one.__class__, RodBase) - or not issubclass(system_two.__class__, RodBase) - or system_one != system_two - ): - raise TypeError( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system and second system should be the same rod \n" - " If you want rod rod contact, use RodRodContact instead".format( - system_one.__class__, system_two.__class__ - ) - ) + common_check_systems_validity(system_one, self._allowed_system_one) + common_check_systems_validity(system_two, self._allowed_system_two) + common_check_systems_different(system_one, system_two) def apply_contact(self, system_one: RodType, system_two: RodType) -> None: """ @@ -390,10 +319,8 @@ def apply_contact(self, system_one: RodType, system_two: RodType) -> None: Parameters ---------- - system_one: object - Rod object. - system_two: object - Rod object. + system_one: RodType + system_two: RodType """ _calculate_contact_forces_self_rod( @@ -439,8 +366,8 @@ def __init__( self, k: np.floating, nu: np.floating, - velocity_damping_coefficient: np.floating = 0.0, - friction_coefficient: np.floating = 0.0, + velocity_damping_coefficient: float = 0.0, + friction_coefficient: float = 0.0, ) -> None: """ Parameters @@ -459,44 +386,20 @@ def __init__( self.k = k self.nu = nu self.velocity_damping_coefficient = velocity_damping_coefficient - self.friction_coefficient = friction_coefficient + self.friction_coefficient = np.float64(friction_coefficient) - def _check_systems_validity( - self, - system_one: SystemType, - system_two: AllowedContactType, - ) -> None: - """ - This checks the contact order and type of a SystemType object and an AllowedContactType object. - For the RodSphereContact class first_system should be a rod and second_system should be a sphere. - Parameters - ---------- - system_one - SystemType - system_two - AllowedContactType - """ - if not issubclass(system_one.__class__, RodBase) or not issubclass( - system_two.__class__, Sphere - ): - raise TypeError( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a sphere".format( - system_one.__class__, system_two.__class__ - ) - ) + @property + def _allowed_system_two(self) -> list[Type]: + return [Sphere] - def apply_contact(self, system_one: RodType, system_two: SystemType) -> None: + def apply_contact(self, system_one: RodType, system_two: Sphere) -> None: """ Apply contact forces and torques between RodType object and Sphere object. Parameters ---------- - system_one: object - Rod object. - system_two: object - Sphere object. + system_one: RodType + system_two: Sphere """ # First, check for a global AABB bounding box, and see whether that @@ -507,7 +410,7 @@ def apply_contact(self, system_one: RodType, system_two: SystemType) -> None: system_one.lengths, system_two.position_collection, system_two.director_collection, - system_two.radius[0], + system_two.radius, ): return @@ -578,33 +481,11 @@ def __init__( self.nu = nu self.surface_tol = 1e-4 - def _check_systems_validity( - self, - system_one: SystemType, - system_two: AllowedContactType, - ) -> None: - """ - This checks the contact order and type of a SystemType object and an AllowedContactType object. - For the RodPlaneContact class first_system should be a rod and second_system should be a plane. - Parameters - ---------- - system_one - SystemType - system_two - AllowedContactType - """ - if not issubclass(system_one.__class__, RodBase) or not issubclass( - system_two.__class__, Plane - ): - raise TypeError( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a plane".format( - system_one.__class__, system_two.__class__ - ) - ) + @property + def _allowed_system_two(self) -> list[Type]: + return [SurfaceBase] - def apply_contact(self, system_one: RodType, system_two: SystemType) -> None: + def apply_contact(self, system_one: RodType, system_two: SurfaceType) -> None: """ Apply contact forces and torques between RodType object and Plane object. @@ -692,42 +573,18 @@ def __init__( self.kinetic_mu_sideways, ) = kinetic_mu_array - def _check_systems_validity( - self, - system_one: SystemType, - system_two: AllowedContactType, - ) -> None: - """ - This checks the contact order and type of a SystemType object and an AllowedContactType object. - For the RodSphereContact class first_system should be a rod and second_system should be a plane. - Parameters - ---------- - system_one - SystemType - system_two - AllowedContactType - """ - if not issubclass(system_one.__class__, RodBase) or not issubclass( - system_two.__class__, Plane - ): - raise TypeError( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a plane".format( - system_one.__class__, system_two.__class__ - ) - ) + @property + def _allowed_system_two(self) -> list[Type]: + return [SurfaceBase] - def apply_contact(self, system_one: RodType, system_two: SystemType) -> None: + def apply_contact(self, system_one: RodType, system_two: SurfaceType) -> None: """ Apply contact forces and torques between RodType object and Plane object with anisotropic friction. Parameters ---------- - system_one: object - Rod object. - system_two: object - Plane object. + system_one: RodType + system_two: SurfaceType """ @@ -794,35 +651,15 @@ def __init__( self.nu = nu self.surface_tol = 1e-4 - def _check_systems_validity( - self, - system_one: SystemType, - system_two: AllowedContactType, - ) -> None: - """ - This checks the contact order and type of a SystemType object and an AllowedContactType object. - For the RodPlaneContact class first_system should be a cylinder and second_system should be a plane. - Parameters - ---------- - system_one - SystemType - system_two - AllowedContactType - """ - if not issubclass(system_one.__class__, Cylinder) or not issubclass( - system_two.__class__, Plane - ): - raise TypeError( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a cylinder, second should be a plane".format( - system_one.__class__, system_two.__class__ - ) - ) + @property + def _allowed_system_one(self) -> list[Type]: + return [Cylinder] - def apply_contact( - self, system_one: Cylinder, system_two: SystemType - ) -> tuple[NDArray[np.floating], NDArray[np.intp]]: + @property + def _allowed_system_two(self) -> list[Type]: + return [SurfaceBase] + + def apply_contact(self, system_one: Cylinder, system_two: SurfaceType) -> None: """ This function computes the plane force response on the cylinder, in the case of contact. Contact model given in Eqn 4.8 Gazzola et. al. RSoS 2018 paper @@ -830,13 +667,11 @@ def apply_contact( Parameters ---------- - system_one: object - Cylinder object. - system_two: object - Plane object. + system_one: Cylinder + system_two: SurfaceBase """ - return _calculate_contact_forces_cylinder_plane( + _calculate_contact_forces_cylinder_plane( system_two.origin, system_two.normal, self.surface_tol, @@ -847,3 +682,49 @@ def apply_contact( system_one.velocity_collection, system_one.external_forces, ) + + +def common_check_systems_identity( + system_one: S1, + system_two: S2, +) -> None: + """ + This checks if two objects are identical. + + Raises + ------ + TypeError + If two objects are identical. + """ + if system_one == system_two: + raise TypeError( + "First system is identical to second system. Systems must be distinct for contact." + ) + + +def common_check_systems_different( + system_one: S1, + system_two: S2, +) -> None: + """ + This checks if two objects are identical. + + Raises + ------ + TypeError + If two objects are not identical. + """ + if system_one != system_two: + raise TypeError("First system must be identical to the second system.") + + +def common_check_systems_validity( + system: S1 | S2, allowed_system: list[Type[S1] | Type[S2]] +) -> None: + # Check validity + if not isinstance(system, tuple(allowed_system)): + system_name = system.__class__.__name__ + allowed_system_names = [candidate.__name__ for candidate in allowed_system] + raise TypeError( + f"System provided ({system_name}) must be derived from {allowed_system_names}." + ) diff --git a/elastica/contact_utils.py b/elastica/contact_utils.py index 657b6dd00..a67793a07 100644 --- a/elastica/contact_utils.py +++ b/elastica/contact_utils.py @@ -7,37 +7,34 @@ from elastica._linalg import ( _batch_norm, ) -from typing import Literal, Sequence, TypeVar +from typing import Literal, Sequence -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _dot_product(a: Sequence[np.floating], b: Sequence[np.floating]) -> np.floating: - sum: np.floating = 0.0 + total: np.float64 = np.float64(0.0) for i in range(3): - sum += a[i] * b[i] - return sum + total += a[i] * b[i] + return total -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _norm(a: Sequence[np.floating]) -> float: return sqrt(_dot_product(a, a)) -_SupportsCompareT = TypeVar("_SupportsCompareT") - - -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _clip(x: np.floating, low: np.floating, high: np.floating) -> np.floating: return max(low, min(x, high)) # Can this be made more efficient than 2 comp, 1 or? -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _out_of_bounds(x: np.floating, low: np.floating, high: np.floating) -> bool: return (x < low) or (x > high) -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _find_min_dist( x1: NDArray[np.floating], e1: NDArray[np.floating], @@ -108,7 +105,7 @@ def _find_min_dist( return x2 + s * e2 - x1 - t * e1, x2 + s * e2, x1 - t * e1 -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _aabbs_not_intersecting( aabb_one: NDArray[np.floating], aabb_two: NDArray[np.floating] ) -> Literal[1, 0]: @@ -123,7 +120,7 @@ def _aabbs_not_intersecting( return 0 -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _prune_using_aabbs_rod_cylinder( rod_one_position_collection: NDArray[np.floating], rod_one_radius_collection: NDArray[np.floating], @@ -165,7 +162,7 @@ def _prune_using_aabbs_rod_cylinder( return _aabbs_not_intersecting(aabb_cylinder, aabb_rod) -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _prune_using_aabbs_rod_rod( rod_one_position_collection: NDArray[np.floating], rod_one_radius_collection: NDArray[np.floating], @@ -203,7 +200,7 @@ def _prune_using_aabbs_rod_rod( return _aabbs_not_intersecting(aabb_rod_two, aabb_rod_one) -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _prune_using_aabbs_rod_sphere( rod_one_position_collection: NDArray[np.floating], rod_one_radius_collection: NDArray[np.floating], @@ -242,7 +239,7 @@ def _prune_using_aabbs_rod_sphere( return _aabbs_not_intersecting(aabb_sphere, aabb_rod) -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _find_slipping_elements( velocity_slip: NDArray[np.floating], velocity_threshold: np.floating ) -> NDArray[np.floating]: @@ -285,7 +282,7 @@ def _find_slipping_elements( return slip_function -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _node_to_element_mass_or_force(input: NDArray[np.floating]) -> NDArray[np.floating]: """ This function converts the mass/forces on rod nodes to @@ -323,7 +320,7 @@ def _node_to_element_mass_or_force(input: NDArray[np.floating]) -> NDArray[np.fl return output -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _elements_to_nodes_inplace( vector_in_element_frame: NDArray[np.floating], vector_in_node_frame: NDArray[np.floating], @@ -349,7 +346,7 @@ def _elements_to_nodes_inplace( vector_in_node_frame[i, k + 1] += 0.5 * vector_in_element_frame[i, k] -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _node_to_element_position( node_position_collection: NDArray[np.floating], ) -> NDArray[np.floating]: @@ -397,7 +394,7 @@ def _node_to_element_position( return element_position_collection -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _node_to_element_velocity( mass: NDArray[np.floating], node_velocity_collection: NDArray[np.floating] ) -> NDArray[np.floating]: diff --git a/elastica/interaction.py b/elastica/interaction.py index f0b9a9b28..d2608bf01 100644 --- a/elastica/interaction.py +++ b/elastica/interaction.py @@ -17,7 +17,7 @@ from numpy.typing import NDArray -from elastica.typing import SystemType +from elastica.typing import SystemType, RodType def find_slipping_elements(velocity_slip: Any, velocity_threshold: Any) -> NoReturn: @@ -110,7 +110,7 @@ def __init__( self.surface_tol = 1e-4 def apply_normal_force( - self, system: SystemType + self, system: RodType ) -> tuple[NDArray[np.floating], NDArray[np.intp]]: """ In the case of contact with the plane, this function computes the plane reaction force on the element. @@ -512,7 +512,7 @@ def __init__(self, dynamic_viscosity: np.floating) -> None: super(SlenderBodyTheory, self).__init__() self.dynamic_viscosity = dynamic_viscosity - def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: + def apply_forces(self, system: RodType, time: np.floating = 0.0) -> None: """ This function applies hydrodynamic forces on body using the slender body theory given in diff --git a/examples/ContinuumSnakeWithLiftingWaveCase/snake_contact.py b/examples/ContinuumSnakeWithLiftingWaveCase/snake_contact.py index 68dea81dd..c27ffd37c 100755 --- a/examples/ContinuumSnakeWithLiftingWaveCase/snake_contact.py +++ b/examples/ContinuumSnakeWithLiftingWaveCase/snake_contact.py @@ -22,10 +22,11 @@ _node_to_element_mass_or_force, ) from numba import njit -from elastica.rod import RodBase +from elastica.rod.rod_base import RodBase from elastica.surface import Plane +from elastica.surface.surface_base import SurfaceBase from elastica.contact_forces import NoContact -from elastica.typing import RodType, SystemType, AllowedContactType +from elastica.typing import RodType, SystemType @njit(cache=True) @@ -286,31 +287,10 @@ def __init__( self.kinetic_mu_sideways, ) = kinetic_mu_array - def _check_systems_validity( - self, - system_one: SystemType, - system_two: AllowedContactType, - ) -> None: - """ - This checks the contact order and type of a SystemType object and an AllowedContactType object. - For the RodSphereContact class first_system should be a rod and second_system should be a plane. - Parameters - ---------- - system_one - SystemType - system_two - AllowedContactType - """ - if not issubclass(system_one.__class__, RodBase) or not issubclass( - system_two.__class__, Plane - ): - raise TypeError( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a plane".format( - system_one.__class__, system_two.__class__ - ) - ) + @property + def _allowed_system_two(self) -> list[Type]: + # Modify this list to include the allowed system types for contact + return [SurfaceBase] def apply_contact(self, system_one: RodType, system_two: SystemType) -> None: """ diff --git a/tests/test_contact_classes.py b/tests/test_contact_classes.py index 94b41a457..509d8e32d 100644 --- a/tests/test_contact_classes.py +++ b/tests/test_contact_classes.py @@ -52,8 +52,8 @@ def mock_cylinder_init(self): self.director_collection = np.array( [[[1.0], [0.0], [0.0]], [[0.0], [1.0], [0.0]], [[0.0], [0.0], [1.0]]] ) - self.radius = np.array([1.0]) - self.length = np.array([2.0]) + self.radius = 1.0 + self.length = 2.0 self.external_forces = np.array([[0.0], [0.0], [0.0]]) self.external_torques = np.array([[0.0], [0.0], [0.0]]) self.velocity_collection = np.array([[0.0], [0.0], [0.0]]) @@ -69,7 +69,7 @@ def mock_sphere_init(self): self.director_collection = np.array( [[[1.0], [0.0], [0.0]], [[0.0], [1.0], [0.0]], [[0.0], [0.0], [1.0]]] ) - self.radius = np.array([1.0]) + self.radius = 1.0 self.velocity_collection = np.array([[0.0], [0.0], [0.0]]) self.external_forces = np.array([[0.0], [0.0], [0.0]]) self.external_torques = np.array([[0.0], [0.0], [0.0]]) @@ -104,43 +104,41 @@ def test_check_systems_validity_with_invalid_systems( "Testing Rod Cylinder Contact wrapper with incorrect type for second argument" with pytest.raises(TypeError) as excinfo: rod_cylinder_contact._check_systems_validity(mock_rod, mock_list) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a cylinder" - ).format(mock_rod.__class__, mock_list.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['Cylinder']." == str( + excinfo.value + ) - "Testing Rod Cylinder Contact wrapper with incorrect type for first argument" with pytest.raises(TypeError) as excinfo: rod_cylinder_contact._check_systems_validity(mock_list, mock_rod) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a cylinder" - ).format(mock_list.__class__, mock_rod.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['RodBase']." == str( + excinfo.value + ) - "Testing Rod Cylinder Contact wrapper with incorrect order" with pytest.raises(TypeError) as excinfo: rod_cylinder_contact._check_systems_validity(mock_cylinder, mock_rod) - print(excinfo.value) assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a cylinder" - ).format(mock_cylinder.__class__, mock_rod.__class__) == str(excinfo.value) + "System provided (MockCylinder) must be derived from ['RodBase']." + == str(excinfo.value) + ) + + with pytest.raises(TypeError) as excinfo: + rod_cylinder_contact._check_systems_validity(mock_list, mock_cylinder) + assert "System provided (list) must be derived from ['RodBase']." == str( + excinfo.value + ) def test_contact_rod_cylinder_with_collision_with_k_without_nu_and_friction( self, ): - "Testing Rod Cylinder Contact wrapper with Collision with analytical verified values" + # Testing Rod Cylinder Contact wrapper with Collision with analytical verified values mock_rod = MockRod() mock_cylinder = MockCylinder() rod_cylinder_contact = RodCylinderContact(k=1.0, nu=0.0) rod_cylinder_contact.apply_contact(mock_rod, mock_cylinder) - """Details and reasoning about the values are given in 'test_contact_specific_functions.py/test_calculate_contact_forces_rod_cylinder()'""" + # Details and reasoning about the values are given in 'test_contact_specific_functions.py/test_calculate_contact_forces_rod_cylinder()' assert_allclose( mock_rod.external_forces, np.array([[0.166666, 0.333333, 0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]), @@ -293,32 +291,27 @@ def test_check_systems_validity_with_invalid_systems( mock_list = [1, 2, 3] rod_rod_contact = RodRodContact(k=1.0, nu=0.0) - "Testing Rod Rod Contact wrapper with incorrect type for second argument" + # Testing Rod Rod Contact wrapper with incorrect type for second argument with pytest.raises(TypeError) as excinfo: rod_rod_contact._check_systems_validity(mock_rod_one, mock_list) - assert ( - "Systems provided to the contact class have incorrect order. \n" - " First system is {0} and second system is {1}. \n" - " Both systems must be distinct rods" - ).format(mock_rod_one.__class__, mock_list.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['RodBase']." == str( + excinfo.value + ) - "Testing Rod Rod Contact wrapper with incorrect type for first argument" + # Testing Rod Rod Contact wrapper with incorrect type for first argument with pytest.raises(TypeError) as excinfo: rod_rod_contact._check_systems_validity(mock_list, mock_rod_one) - assert ( - "Systems provided to the contact class have incorrect order. \n" - " First system is {0} and second system is {1}. \n" - " Both systems must be distinct rods" - ).format(mock_list.__class__, mock_rod_one.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['RodBase']." == str( + excinfo.value + ) - "Testing Rod Rod Contact wrapper with same rod for both arguments" + # Testing Rod Rod Contact wrapper with same rod for both arguments with pytest.raises(TypeError) as excinfo: rod_rod_contact._check_systems_validity(mock_rod_one, mock_rod_one) assert ( - "First rod is identical to second rod. \n" - "Rods must be distinct for RodRodConact. \n" - "If you want self contact, use RodSelfContact instead" - ) == str(excinfo.value) + "First system is identical to second system. Systems must be distinct for contact." + == str(excinfo.value) + ) def test_contact_with_two_rods_with_collision_with_k_without_nu(self): @@ -439,35 +432,26 @@ def test_check_systems_validity_with_invalid_systems( mock_list = [1, 2, 3] self_contact = RodSelfContact(k=1.0, nu=0.0) - "Testing Self Contact wrapper with incorrect type for second argument" + # Testing Self Contact wrapper with incorrect type for second argument with pytest.raises(TypeError) as excinfo: self_contact._check_systems_validity(mock_rod_one, mock_list) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system and second system should be the same rod \n" - " If you want rod rod contact, use RodRodContact instead" - ).format(mock_rod_one.__class__, mock_list.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['RodBase']." == str( + excinfo.value + ) - "Testing Self Contact wrapper with incorrect type for first argument" + # Testing Self Contact wrapper with incorrect type for first argument with pytest.raises(TypeError) as excinfo: self_contact._check_systems_validity(mock_list, mock_rod_one) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system and second system should be the same rod \n" - " If you want rod rod contact, use RodRodContact instead" - ).format(mock_list.__class__, mock_rod_one.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['RodBase']." == str( + excinfo.value + ) - "Testing Self Contact wrapper with different rods" + # Testing Self Contact wrapper with different rods with pytest.raises(TypeError) as excinfo: self_contact._check_systems_validity(mock_rod_one, mock_rod_two) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system and second system should be the same rod \n" - " If you want rod rod contact, use RodRodContact instead" - ).format(mock_rod_one.__class__, mock_rod_two.__class__) == str(excinfo.value) + assert "First system must be identical to the second system." == str( + excinfo.value + ) def test_self_contact_with_rod_self_collision(self): @@ -530,46 +514,39 @@ def test_check_systems_validity_with_invalid_systems( mock_sphere = MockSphere() rod_sphere_contact = RodSphereContact(k=1.0, nu=0.0) - "Testing Rod Sphere Contact wrapper with incorrect type for second argument" + # Testing Rod Sphere Contact wrapper with incorrect type for second argument with pytest.raises(TypeError) as excinfo: rod_sphere_contact._check_systems_validity(mock_rod, mock_list) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a sphere" - ).format(mock_rod.__class__, mock_list.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['Sphere']." == str( + excinfo.value + ) - "Testing Rod Sphere Contact wrapper with incorrect type for first argument" + # Testing Rod Sphere Contact wrapper with incorrect type for first argument with pytest.raises(TypeError) as excinfo: rod_sphere_contact._check_systems_validity(mock_list, mock_rod) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a sphere" - ).format(mock_list.__class__, mock_rod.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['RodBase']." == str( + excinfo.value + ) - "Testing Rod Sphere Contact wrapper with incorrect order" + # Testing Rod Sphere Contact wrapper with incorrect order with pytest.raises(TypeError) as excinfo: rod_sphere_contact._check_systems_validity(mock_sphere, mock_rod) - print(excinfo.value) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a sphere" - ).format(mock_sphere.__class__, mock_rod.__class__) == str(excinfo.value) + assert "System provided (MockSphere) must be derived from ['RodBase']." == str( + excinfo.value + ) def test_contact_rod_sphere_with_collision_with_k_without_nu_and_friction( self, ): - "Testing Rod Sphere Contact wrapper with Collision with analytical verified values" + # "Testing Rod Sphere Contact wrapper with Collision with analytical verified values mock_rod = MockRod() mock_sphere = MockSphere() rod_sphere_contact = RodSphereContact(k=1.0, nu=0.0) rod_sphere_contact.apply_contact(mock_rod, mock_sphere) - """Details and reasoning about the values are given in 'test_contact_specific_functions.py/test_calculate_contact_forces_rod_sphere_with_k_without_nu_and_friction()'""" + # Details and reasoning about the values are given in 'test_contact_specific_functions.py/test_calculate_contact_forces_rod_sphere_with_k_without_nu_and_friction()' assert_allclose( mock_sphere.external_forces, np.array([[-0.5], [0], [0]]), atol=1e-6 ) @@ -630,23 +607,19 @@ def test_check_systems_validity_with_invalid_systems( mock_list = [1, 2, 3] rod_plane_contact = RodPlaneContact(k=1.0, nu=0.0) - "Testing Rod Plane Contact wrapper with incorrect type for second argument" + # Testing Rod Plane Contact wrapper with incorrect type for second argument with pytest.raises(TypeError) as excinfo: rod_plane_contact._check_systems_validity(mock_rod, mock_list) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a plane" - ).format(mock_rod.__class__, mock_list.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['SurfaceBase']." == str( + excinfo.value + ) - "Testing Rod Plane Contact wrapper with incorrect type for first argument" + # Testing Rod Plane Contact wrapper with incorrect type for first argument with pytest.raises(TypeError) as excinfo: rod_plane_contact._check_systems_validity(mock_list, mock_plane) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a plane" - ).format(mock_list.__class__, mock_plane.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['RodBase']." == str( + excinfo.value + ) def test_rod_plane_contact_without_contact(self): """ @@ -905,23 +878,19 @@ def test_check_systems_validity_with_invalid_systems( np.array([0.0, 0.0, 0.0]), # forward, backward, sideways ) - "Testing Rod Plane Contact wrapper with incorrect type for second argument" + # Testing Rod Plane Contact wrapper with incorrect type for second argument with pytest.raises(TypeError) as excinfo: rod_plane_contact._check_systems_validity(mock_rod, mock_list) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a plane" - ).format(mock_rod.__class__, mock_list.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['SurfaceBase']." == str( + excinfo.value + ) - "Testing Rod Plane wrapper with incorrect type for first argument" + # Testing Rod Plane wrapper with incorrect type for first argument with pytest.raises(TypeError) as excinfo: rod_plane_contact._check_systems_validity(mock_list, mock_plane) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a rod, second should be a plane" - ).format(mock_list.__class__, mock_plane.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['RodBase']." == str( + excinfo.value + ) @pytest.mark.parametrize("velocity", [-1.0, -3.0, 1.0, 5.0, 2.0]) def test_axial_kinetic_friction(self, velocity): @@ -1212,7 +1181,7 @@ def initializer( # create plane plane = MockPlane() - plane.origin = np.array([0.0, -cylinder.radius[0] + shift, 0.0]).reshape(3, 1) + plane.origin = np.array([0.0, -cylinder.radius + shift, 0.0]).reshape(3, 1) plane.normal = plane_normal.reshape( 3, ) @@ -1232,23 +1201,19 @@ def test_check_systems_validity_with_invalid_systems( mock_list = [1, 2, 3] cylinder_plane_contact = CylinderPlaneContact(k=1.0, nu=0.0) - "Testing Cylinder Plane Contact wrapper with incorrect type for second argument" + # Testing Cylinder Plane Contact wrapper with incorrect type for second argument with pytest.raises(TypeError) as excinfo: cylinder_plane_contact._check_systems_validity(mock_cylinder, mock_list) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a cylinder, second should be a plane" - ).format(mock_cylinder.__class__, mock_list.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['SurfaceBase']." == str( + excinfo.value + ) - "Testing Cylinder Plane wrapper with incorrect type for first argument" + # Testing Cylinder Plane wrapper with incorrect type for first argument with pytest.raises(TypeError) as excinfo: cylinder_plane_contact._check_systems_validity(mock_list, mock_plane) - assert ( - "Systems provided to the contact class have incorrect order/type. \n" - " First system is {0} and second system is {1}. \n" - " First system should be a cylinder, second should be a plane" - ).format(mock_list.__class__, mock_plane.__class__) == str(excinfo.value) + assert "System provided (list) must be derived from ['Cylinder']." == str( + excinfo.value + ) def test_cylinder_plane_contact_without_contact(self): """ diff --git a/tests/test_modules/test_contact.py b/tests/test_modules/test_contact.py index 82c41f696..5b9ffe75b 100644 --- a/tests/test_modules/test_contact.py +++ b/tests/test_modules/test_contact.py @@ -240,9 +240,17 @@ def mock_init(self, *args, **kwargs): pass # in place class - MockContact = type( - "MockContact", (self.NoContact, object), {"__init__": mock_init} - ) + class MockContact(self.NoContact): + def __init__(self, *args, **kwargs): + pass + + @property + def _allowed_system_one(self): + return [TestContactMixin.MockRod] + + @property + def _allowed_system_two(self): + return [TestContactMixin.MockRod] # Constrain any and all systems system_collection_with_contacts.detect_contact_between(0, 1).using( @@ -283,9 +291,17 @@ def mock_init(self, *args, **kwargs): pass # in place class - MockContact = type( - "MockContact", (self.NoContact, object), {"__init__": mock_init} - ) + class MockContact(self.NoContact): + def __init__(self, *args, **kwargs): + pass + + @property + def _allowed_system_one(self): + return [TestContactMixin.MockRod] + + @property + def _allowed_system_two(self): + return [TestContactMixin.MockRigidBody] # incorrect order contact system_collection_with_contacts.detect_contact_between( @@ -302,17 +318,12 @@ def test_contact_check_order(self, load_contact_objects_with_incorrect_order): contact_cls, ) = load_contact_objects_with_incorrect_order - mock_rod = self.MockRod(2, 3, 4, 5) - mock_rigid_body = self.MockRigidBody(5.0, 5.0) - with pytest.raises(TypeError) as excinfo: system_collection_with_contacts._finalize_contact() assert ( - "Systems provided to the contact class have incorrect order. \n" - " First system is {0} and second system is {1}. \n" - " If the first system is a rod, the second system can be a rod, rigid body or surface. \n" - " If the first system is a rigid body, the second system can be a rigid body or surface." - ).format(mock_rigid_body.__class__, mock_rod.__class__) == str(excinfo.value) + "System provided (MockRigidBody) must be derived from ['MockRod']" + in str(excinfo.value) + ) @pytest.fixture def load_system_with_rods_in_contact(self, load_system_with_contacts): From 8a056a660167e33bc660f9a3eb355e60b31c0eb5 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 16 Jun 2024 00:50:17 +0900 Subject: [PATCH 085/134] remove AllowedContactType union alias to distinguish cases --- elastica/modules/base_system.py | 1 - elastica/modules/contact.py | 7 ++----- elastica/modules/protocol.py | 7 ++----- elastica/typing.py | 11 +++++------ 4 files changed, 9 insertions(+), 17 deletions(-) diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index fd999665f..8e91fd1b3 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -13,7 +13,6 @@ OperatorType, OperatorCallbackType, OperatorFinalizeType, - AllowedContactType, ) import numpy as np diff --git a/elastica/modules/contact.py b/elastica/modules/contact.py index 924e7a4a9..69cb11db0 100644 --- a/elastica/modules/contact.py +++ b/elastica/modules/contact.py @@ -11,7 +11,6 @@ SystemIdxType, OperatorFinalizeType, SystemType, - AllowedContactType, ) from .protocol import SystemCollectionProtocol, ModuleProtocol @@ -48,7 +47,7 @@ def __init__(self: SystemCollectionProtocol) -> None: def detect_contact_between( self: SystemCollectionProtocol, first_system: SystemType, - second_system: AllowedContactType, + second_system: SystemType, ) -> ModuleProtocol: """ This method adds contact detection between two objects using the selected contact class. @@ -57,9 +56,7 @@ def detect_contact_between( Parameters ---------- first_system : SystemType - Rod or rigid body object - second_system : AllowedContactType - Rod, rigid body or surface object + second_system : SystemType Returns ------- diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index 5ce5ba7cf..f27544650 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -7,7 +7,6 @@ OperatorCallbackType, OperatorFinalizeType, SystemType, - AllowedContactType, ConnectionIndex, ) from elastica.joint import FreeJoint @@ -59,9 +58,7 @@ def _feature_group_finalize(self) -> list[OperatorFinalizeType]: ... def blocks(self) -> Generator[SystemType, None, None]: ... - def _get_sys_idx_if_valid( - self, sys_to_be_added: SystemType | AllowedContactType - ) -> SystemIdxType: ... + def _get_sys_idx_if_valid(self, sys_to_be_added: SystemType) -> SystemIdxType: ... # Connection API _finalize_connections: OperatorFinalizeType @@ -123,7 +120,7 @@ def add_forcing_to(self, system: SystemType) -> ModuleProtocol: @abstractmethod def detect_contact_between( - self, first_system: SystemType, second_system: AllowedContactType + self, first_system: SystemType, second_system: SystemType ) -> ModuleProtocol: raise NotImplementedError diff --git a/elastica/typing.py b/elastica/typing.py index 43ccb3df8..b7b51eb29 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -4,7 +4,7 @@ """ from typing import TYPE_CHECKING -from typing import Type, Callable, Any, ParamSpec +from typing import Callable, Any, ParamSpec from typing import TypeAlias import numpy as np @@ -14,8 +14,8 @@ # Used for type hinting without circular imports # NEVER BACK-IMPORT ANY ELASTICA MODULES HERE from .rod.protocol import CosseratRodProtocol - from .rigidbody import RigidBodyBase - from .surface import SurfaceBase + from .rigidbody.rigid_body import RigidBodyBase + from .surface.surface_base import SurfaceBase from .modules.base_system import BaseSystemCollection from .modules.protocol import SystemCollectionProtocol @@ -56,10 +56,9 @@ RodType: TypeAlias = "CosseratRodProtocol" RigidBodyType: TypeAlias = "RigidBodyProtocol" +SurfaceType: TypeAlias = "SurfaceBase" + SystemCollectionType: TypeAlias = "SystemCollectionProtocol" -AllowedContactType: TypeAlias = ( - SystemType | Type["SurfaceBase"] -) # FIXME: SurfaceBase needs to be treated differently BlockType: TypeAlias = "BlockCosseratRodProtocol" # | "BlockRigidBodyProtocol" # Indexing types From 74845025209d16ce456cf419adaf2ff5ec36bcc1 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 16 Jun 2024 00:50:50 +0900 Subject: [PATCH 086/134] typing: add basic properties for surface_base definition --- elastica/restart.py | 2 +- elastica/surface/surface_base.py | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/elastica/restart.py b/elastica/restart.py index 643e0ed03..cd31d520f 100644 --- a/elastica/restart.py +++ b/elastica/restart.py @@ -30,7 +30,7 @@ def all_equal(iterable: Iterable[Any]) -> bool: def save_state( simulator: Iterable, directory: str = "", - time: np.floating = 0.0, + time: np.floating = np.float64(0.0), verbose: bool = False, ) -> None: """ diff --git a/elastica/surface/surface_base.py b/elastica/surface/surface_base.py index 421cb1304..c5c7a2ff7 100644 --- a/elastica/surface/surface_base.py +++ b/elastica/surface/surface_base.py @@ -1,6 +1,9 @@ __doc__ = """Base class for surfaces""" from typing import Type +import numpy as np +from numpy.typing import NDArray + class SurfaceBase: """ @@ -18,4 +21,5 @@ def __init__(self) -> None: """ SurfaceBase does not take any arguments. """ - pass + self.normal: NDArray[np.floating] # (3,) + self.origin: NDArray[np.floating] # (3, 1) From 64e39a661d01e5c543e0447bbe151e888a987b20 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 16 Jun 2024 01:10:15 +0900 Subject: [PATCH 087/134] typing: include interaction module --- elastica/external_forces.py | 10 +++++++--- elastica/interaction.py | 18 +++++++++++------- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/elastica/external_forces.py b/elastica/external_forces.py index 00cbe60d0..faabf4fc7 100644 --- a/elastica/external_forces.py +++ b/elastica/external_forces.py @@ -1,6 +1,7 @@ __doc__ = """ Numba implementation module for boundary condition implementations that apply external forces to the system.""" +from typing import TypeVar, Generic import numpy as np from numpy.typing import NDArray @@ -13,7 +14,10 @@ from elastica._linalg import _batch_product_i_k_to_ik -class NoForces: +S = TypeVar("S") + + +class NoForces(Generic[S]): """ This is the base class for external forcing boundary conditions applied to rod-like objects. @@ -30,7 +34,7 @@ def __init__(self) -> None: """ pass - def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: + def apply_forces(self, system: S, time: np.floating = np.float64(0.0)) -> None: """Apply forces to a rod-like object. In NoForces class, this routine simply passes. @@ -45,7 +49,7 @@ def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: """ pass - def apply_torques(self, system: SystemType, time: np.floating = 0.0) -> None: + def apply_torques(self, system: S, time: np.floating = np.float64(0.0)) -> None: """Apply torques to a rod-like object. In NoForces class, this routine simply passes. diff --git a/elastica/interaction.py b/elastica/interaction.py index d2608bf01..d5ec64c33 100644 --- a/elastica/interaction.py +++ b/elastica/interaction.py @@ -17,7 +17,7 @@ from numpy.typing import NDArray -from elastica.typing import SystemType, RodType +from elastica.typing import SystemType, RodType, RigidBodyType def find_slipping_elements(velocity_slip: Any, velocity_threshold: Any) -> NoReturn: @@ -247,12 +247,14 @@ def __init__( # kinetic and static friction should separate functions # for now putting them together to figure out common variables - def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: + def apply_forces( + self, system: RodType | RigidBodyType, time: np.floating = np.float64(0) + ) -> None: """ Call numba implementation to apply friction forces Parameters ---------- - system + system : RodType | RigidBodyType time """ @@ -349,7 +351,7 @@ def sum_over_elements(input: NDArray[np.floating]) -> np.floating: This version: 513 ns ± 24.6 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each) """ - output: np.floating = 0.0 + output: np.floating = np.float64(0.0) for i in range(input.shape[0]): output += input[i] @@ -512,7 +514,9 @@ def __init__(self, dynamic_viscosity: np.floating) -> None: super(SlenderBodyTheory, self).__init__() self.dynamic_viscosity = dynamic_viscosity - def apply_forces(self, system: RodType, time: np.floating = 0.0) -> None: + def apply_forces( + self, system: RodType, time: np.floating = np.float64(0.0) + ) -> None: """ This function applies hydrodynamic forces on body using the slender body theory given in @@ -537,7 +541,7 @@ def apply_forces(self, system: RodType, time: np.floating = 0.0) -> None: # base class for interaction # only applies normal force no friction -class InteractionPlaneRigidBody: +class InteractionPlaneRigidBody(InteractionPlane): def __init__( self, k: np.floating, @@ -552,7 +556,7 @@ def __init__( self.surface_tol = 1e-4 def apply_normal_force( - self, system: SystemType + self, system: RigidBodyType ) -> tuple[NDArray[np.floating], NDArray[np.intp]]: """ This function computes the plane force response on the rigid body, in the From 32130fbe786c0782e02109a10df42db2d9a0ee84 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 16 Jun 2024 10:05:32 +0900 Subject: [PATCH 088/134] typing: other minor update for different operators --- elastica/_synchronize_periodic_boundary.py | 6 ++-- elastica/boundary_conditions.py | 19 +++++----- elastica/callback_functions.py | 11 +++--- elastica/dissipation.py | 10 +++--- elastica/external_forces.py | 40 ++++++++++++++-------- elastica/rod/protocol.py | 2 ++ 6 files changed, 53 insertions(+), 35 deletions(-) diff --git a/elastica/_synchronize_periodic_boundary.py b/elastica/_synchronize_periodic_boundary.py index 7d4c4882b..49784f6c6 100644 --- a/elastica/_synchronize_periodic_boundary.py +++ b/elastica/_synchronize_periodic_boundary.py @@ -7,7 +7,7 @@ import numpy as np from numpy.typing import NDArray from elastica.boundary_conditions import ConstraintBase -from elastica.typing import SystemType +from elastica.typing import RodType @njit(cache=True) # type: ignore @@ -92,7 +92,7 @@ class _ConstrainPeriodicBoundaries(ConstraintBase): def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) - def constrain_values(self, rod: SystemType, time: np.floating) -> None: + def constrain_values(self, rod: RodType, time: np.floating) -> None: _synchronize_periodic_boundary_of_vector_collection( rod.position_collection, rod.periodic_boundary_nodes_idx ) @@ -100,7 +100,7 @@ def constrain_values(self, rod: SystemType, time: np.floating) -> None: rod.director_collection, rod.periodic_boundary_elems_idx ) - def constrain_rates(self, rod: SystemType, time: np.floating) -> None: + def constrain_rates(self, rod: RodType, time: np.floating) -> None: _synchronize_periodic_boundary_of_vector_collection( rod.velocity_collection, rod.periodic_boundary_nodes_idx ) diff --git a/elastica/boundary_conditions.py b/elastica/boundary_conditions.py index 3a1b99642..d351ed323 100644 --- a/elastica/boundary_conditions.py +++ b/elastica/boundary_conditions.py @@ -1,7 +1,7 @@ __doc__ = """ Built-in boundary condition implementationss """ import warnings -from typing import Any, Optional +from typing import Any, Optional, TypeVar, Generic import numpy as np from numpy.typing import NDArray @@ -15,7 +15,10 @@ from elastica.typing import SystemType, RodType -class ConstraintBase(ABC): +S = TypeVar("S") + + +class ConstraintBase(ABC, Generic[S]): """Base class for constraint and displacement boundary condition implementation. Notes @@ -31,7 +34,7 @@ class ConstraintBase(ABC): """ - _system: SystemType + _system: S _constrained_position_idx: np.ndarray _constrained_director_idx: np.ndarray @@ -51,24 +54,24 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: ) @property - def system(self) -> SystemType: + def system(self) -> S: """get system (rod or rigid body) reference""" return self._system @property - def constrained_position_idx(self) -> Optional[np.ndarray]: + def constrained_position_idx(self) -> np.ndarray: """get position-indices passed to "using" """ # TODO: This should be immutable somehow return self._constrained_position_idx @property - def constrained_director_idx(self) -> Optional[np.ndarray]: + def constrained_director_idx(self) -> np.ndarray: """get director-indices passed to "using" """ # TODO: This should be immutable somehow return self._constrained_director_idx @abstractmethod - def constrain_values(self, system: SystemType, time: np.floating) -> None: + def constrain_values(self, system: S, time: np.floating) -> None: # TODO: In the future, we can remove rod and use self.system """ Constrain values (position and/or directors) of a rod object. @@ -83,7 +86,7 @@ def constrain_values(self, system: SystemType, time: np.floating) -> None: pass @abstractmethod - def constrain_rates(self, system: SystemType, time: np.floating) -> None: + def constrain_rates(self, system: S, time: np.floating) -> None: # TODO: In the future, we can remove rod and use self.system """ Constrain rates (velocity and/or omega) of a rod object. diff --git a/elastica/callback_functions.py b/elastica/callback_functions.py index dac062080..88b15788d 100644 --- a/elastica/callback_functions.py +++ b/elastica/callback_functions.py @@ -1,11 +1,11 @@ __doc__ = """ Module contains callback classes to save simulation data for rod-like objects """ +from typing import Any, Optional, TypeVar, Generic from elastica.typing import SystemType import os import sys import numpy as np import logging -from typing import Any, Optional from collections import defaultdict @@ -13,7 +13,10 @@ from elastica.typing import RodType, SystemType -class CallBackBaseClass: +T = TypeVar("T") + + +class CallBackBaseClass(Generic[T]): """ This is the base class for callbacks for rod-like objects. @@ -30,9 +33,7 @@ def __init__(self) -> None: """ pass - def make_callback( - self, system: SystemType, time: np.floating, current_step: int - ) -> None: + def make_callback(self, system: T, time: np.floating, current_step: int) -> None: """ This method is called every time step. Users can define which parameters are called back and recorded. Also users diff --git a/elastica/dissipation.py b/elastica/dissipation.py index dfe5373af..95453beaf 100644 --- a/elastica/dissipation.py +++ b/elastica/dissipation.py @@ -7,7 +7,7 @@ from abc import ABC, abstractmethod from typing import Any -from elastica.typing import RodType, SystemType +from elastica.typing import RodType, RigidBodyType, SystemType from numba import njit @@ -29,7 +29,7 @@ class DamperBase(ABC): """ - _system: SystemType + _system: RodType | RigidBodyType # TODO typing can be made better def __init__(self, *args: Any, **kwargs: Any) -> None: @@ -43,7 +43,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: ) @property - def system(self) -> SystemType: + def system(self) -> RodType | RigidBodyType: """ get system (rod or rigid body) reference @@ -55,7 +55,7 @@ def system(self) -> SystemType: return self._system @abstractmethod - def dampen_rates(self, system: SystemType, time: np.floating) -> None: + def dampen_rates(self, system: RodType | RigidBodyType, time: np.floating) -> None: # TODO: In the future, we can remove rod and use self.system """ Dampen rates (velocity and/or omega) of a rod object. @@ -148,7 +148,7 @@ def __init__( * np.diagonal(self._system.inv_mass_second_moment_of_inertia).T ) - def dampen_rates(self, rod: SystemType, time: np.floating) -> None: + def dampen_rates(self, rod: RodType | RigidBodyType, time: np.floating) -> None: rod.velocity_collection[:] = ( rod.velocity_collection * self.translational_damping_coefficient ) diff --git a/elastica/external_forces.py b/elastica/external_forces.py index faabf4fc7..f3d3655d7 100644 --- a/elastica/external_forces.py +++ b/elastica/external_forces.py @@ -7,7 +7,7 @@ from numpy.typing import NDArray from elastica._linalg import _batch_matvec -from elastica.typing import SystemType, RodType +from elastica.typing import SystemType, RodType, RigidBodyType from elastica.utils import _bspline from numba import njit @@ -90,7 +90,9 @@ def __init__( super(GravityForces, self).__init__() self.acc_gravity = acc_gravity - def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: + def apply_forces( + self, system: RodType | RigidBodyType, time: np.floating = np.float64(0.0) + ) -> None: self.compute_gravity_forces( self.acc_gravity, system.mass, system.external_forces ) @@ -138,7 +140,7 @@ def __init__( self, start_force: NDArray[np.floating], end_force: NDArray[np.floating], - ramp_up_time: np.floating, + ramp_up_time: float, ) -> None: """ @@ -160,7 +162,9 @@ def __init__( assert ramp_up_time > 0.0 self.ramp_up_time = ramp_up_time - def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: + def apply_forces( + self, system: SystemType, time: np.floating = np.float64(0.0) + ) -> None: self.compute_end_point_forces( system.external_forces, self.start_force, @@ -195,9 +199,9 @@ def compute_end_point_forces( Applied forces are ramped up until ramp up time. """ - factor: np.floating = min(1.0, time / ramp_up_time) - external_forces[..., 0] += start_force * factor - external_forces[..., -1] += end_force * factor + factor = min(1.0, time / ramp_up_time) + external_forces[..., 0] += start_force * factor # type: ignore[operator] + external_forces[..., -1] += end_force * factor # type: ignore[operator] class UniformTorques(NoForces): @@ -229,7 +233,9 @@ def __init__( super(UniformTorques, self).__init__() self.torque = torque * direction - def apply_torques(self, system: SystemType, time: np.floating = 0.0) -> None: + def apply_torques( + self, system: SystemType, time: np.floating = np.float64(0.0) + ) -> None: n_elems = system.n_elems torque_on_one_element = ( _batch_product_i_k_to_ik(self.torque, np.ones((n_elems))) / n_elems @@ -267,7 +273,9 @@ def __init__( super(UniformForces, self).__init__() self.force = (force * direction).reshape(3, 1) - def apply_forces(self, rod: SystemType, time: np.floating = 0.0) -> None: + def apply_forces( + self, rod: SystemType, time: np.floating = np.float64(0.0) + ) -> None: force_on_one_element = self.force / rod.n_elems rod.external_forces += force_on_one_element @@ -331,7 +339,7 @@ def __init__( Phase shift of traveling wave. direction: numpy.ndarray 1D (dim) array containing data with 'float' type. Muscle torque direction. - ramp_up_time: float + ramp_up_time: np.floating Applied muscle torques are ramped up until ramp up time. with_spline: boolean Option to use beta-spline. @@ -364,7 +372,9 @@ def __init__( else: self.my_spline = np.full_like(self.s, fill_value=1.0) - def apply_torques(self, rod: SystemType, time: np.floating = 0.0) -> None: + def apply_torques( + self, rod: SystemType, time: np.floating = np.float64(0.0) + ) -> None: self.compute_muscle_torques( time, self.my_spline, @@ -393,7 +403,7 @@ def compute_muscle_torques( external_torques: NDArray[np.floating], ) -> None: # Ramp up the muscle torque - factor: np.floating = min(1.0, time / ramp_up_time) + factor = np.float64(min(np.float64(1.0), time / ramp_up_time)) # From the node 1 to node nelem-1 # Magnitude of the torque. Am = beta(s) * sin(2pi*t/T + 2pi*s/lambda + phi) # There is an inconsistency with paper and Elastica cpp implementation. In paper sign in @@ -497,7 +507,7 @@ def __init__( self, start_force_mag: np.floating, end_force_mag: np.floating, - ramp_up_time: np.floating = 0.0, + ramp_up_time: float = 0.0, tangent_direction: NDArray[np.floating] = np.array([0, 0, 1]), normal_direction: NDArray[np.floating] = np.array([0, 1, 0]), ) -> None: @@ -530,7 +540,9 @@ def __init__( assert ramp_up_time >= 0.0 self.ramp_up_time = ramp_up_time - def apply_forces(self, system: SystemType, time: np.floating = 0.0) -> None: + def apply_forces( + self, system: SystemType, time: np.floating = np.float64(0.0) + ) -> None: if time < self.ramp_up_time: # When time smaller than ramp up time apply the force in normal direction diff --git a/elastica/rod/protocol.py b/elastica/rod/protocol.py index ba83bb225..672b5a33e 100644 --- a/elastica/rod/protocol.py +++ b/elastica/rod/protocol.py @@ -44,3 +44,5 @@ class CosseratRodProtocol(SystemProtocol, _RodEnergy, Protocol): ghost_elems_idx: NDArray[np.integer] ring_rod_flag: bool + periodic_boundary_nodes_idx: NDArray[np.integer] + periodic_boundary_elems_idx: NDArray[np.integer] From 3a06dc9c47d845704def20f2978e9a688d9271d7 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 16 Jun 2024 10:21:11 +0900 Subject: [PATCH 089/134] remove deprecated joint functions --- elastica/__init__.py | 2 - elastica/joint.py | 279 +++++-------------------------------------- 2 files changed, 27 insertions(+), 254 deletions(-) diff --git a/elastica/__init__.py b/elastica/__init__.py index 7cdfc7f39..f705615a4 100644 --- a/elastica/__init__.py +++ b/elastica/__init__.py @@ -36,10 +36,8 @@ ) from elastica.joint import ( FreeJoint, - ExternalContact, FixedJoint, HingeJoint, - SelfContact, ) from elastica.contact_forces import ( NoContact, diff --git a/elastica/joint.py b/elastica/joint.py index eb3fd5d18..6df1e47ea 100644 --- a/elastica/joint.py +++ b/elastica/joint.py @@ -1,11 +1,13 @@ __doc__ = """ Module containing joint classes to connect multiple rods together. """ +__all__ = ["FreeJoint", "HingeJoint", "FixedJoint", "get_relative_rotation_two_systems"] from typing import Optional from elastica._rotations import _inv_rotate from elastica.typing import SystemType, RodType, ConnectionIndex import numpy as np -import logging + +pass from numpy.typing import NDArray @@ -47,9 +49,9 @@ def __init__(self, k: np.floating, nu: np.floating) -> None: def apply_forces( self, - system_one: SystemType, + system_one: RodType | RigidBodyType, index_one: ConnectionIndex, - system_two: SystemType, + system_two: RodType | RigidBodyType, index_two: ConnectionIndex, ) -> None: """ @@ -57,11 +59,11 @@ def apply_forces( Parameters ---------- - system_one : SystemType + system_one : RodType | RigidBodyType Rod or rigid-body object index_one : ConnectionIndex Index of first rod for joint. - system_two : SystemType + system_two : RodType | RigidBodyType Rod or rigid-body object index_two : ConnectionIndex Index of second rod for joint. @@ -90,9 +92,9 @@ def apply_forces( def apply_torques( self, - system_one: SystemType, + system_one: RodType | RigidBodyType, index_one: ConnectionIndex, - system_two: SystemType, + system_two: RodType | RigidBodyType, index_two: ConnectionIndex, ) -> None: """ @@ -102,11 +104,11 @@ def apply_torques( Parameters ---------- - system_one : SystemType + system_one : RodType | RigidBodyType Rod or rigid-body object index_one : ConnectionIndex Index of first rod for joint. - system_two : SystemType + system_two : RodType | RigidBodyType Rod or rigid-body object index_two : ConnectionIndex Index of second rod for joint. @@ -171,18 +173,18 @@ def __init__( # Apply force is same as free joint def apply_forces( self, - system_one: SystemType, + system_one: RodType | RigidBodyType, index_one: ConnectionIndex, - system_two: SystemType, + system_two: RodType | RigidBodyType, index_two: ConnectionIndex, ) -> None: return super().apply_forces(system_one, index_one, system_two, index_two) def apply_torques( self, - system_one: SystemType, + system_one: RodType | RigidBodyType, index_one: ConnectionIndex, - system_two: SystemType, + system_two: RodType | RigidBodyType, index_two: ConnectionIndex, ) -> None: # current tangent direction of the `index_two` element of system two @@ -278,18 +280,18 @@ def __init__( # Apply force is same as free joint def apply_forces( self, - system_one: SystemType, + system_one: RodType | RigidBodyType, index_one: ConnectionIndex, - system_two: SystemType, + system_two: RodType | RigidBodyType, index_two: ConnectionIndex, ) -> None: return super().apply_forces(system_one, index_one, system_two, index_two) def apply_torques( self, - system_one: SystemType, + system_one: RodType | RigidBodyType, index_one: ConnectionIndex, - system_two: SystemType, + system_two: RodType | RigidBodyType, index_two: ConnectionIndex, ) -> None: # collect directors of systems one and two @@ -335,10 +337,10 @@ def apply_torques( def get_relative_rotation_two_systems( - system_one: SystemType, - index_one: int, - system_two: SystemType, - index_two: int, + system_one: RodType | RigidBodyType, + index_one: ConnectionIndex, + system_two: RodType | RigidBodyType, + index_two: ConnectionIndex, ) -> NDArray[np.floating]: """ Compute the relative rotation matrix C_12 between system one and system two at the specified elements. @@ -366,13 +368,13 @@ def get_relative_rotation_two_systems( Parameters ---------- - system_one : SystemType + system_one : RodType | RigidBodyType Rod or rigid-body object - index_one : int + index_one : ConnectionIndex Index of first rod for joint. - system_two : SystemType + system_two : RodType | RigidBodyType Rod or rigid-body object - index_two : int + index_two : ConnectionIndex Index of second rod for joint. Returns @@ -384,230 +386,3 @@ def get_relative_rotation_two_systems( system_one.director_collection[..., index_one] @ system_two.director_collection[..., index_two].T ) - - -class ExternalContact(FreeJoint): - """ - This class is for applying contact forces between rod-cylinder and rod-rod. - If you are want to apply contact forces between rod and cylinder, first system is always rod and second system - is always cylinder. - In addition to the contact forces, user can define apply friction forces between rod and cylinder that - are in contact. For details on friction model refer to this [1]_. - TODO: Currently friction force is between rod-cylinder, in future implement friction forces between rod-rod. - - Notes - ----- - The `velocity_damping_coefficient` is set to a high value (e.g. 1e4) to minimize slip and simulate stiction - (static friction), while friction_coefficient corresponds to the Coulombic friction coefficient. - - Examples - -------- - How to define contact between rod and cylinder. - - >>> simulator.connect(rod, cylinder).using( - ... ExternalContact, - ... k=1e4, - ... nu=10, - ... velocity_damping_coefficient=10, - ... kinetic_friction_coefficient=10, - ... ) - - How to define contact between rod and rod. - - >>> simulator.connect(rod, rod).using( - ... ExternalContact, - ... k=1e4, - ... nu=10, - ... ) - - .. [1] Preclik T., Popa Constantin., Rude U., Regularizing a Time-Stepping Method for Rigid Multibody Dynamics, Multibody Dynamics 2011, ECCOMAS. URL: https://www10.cs.fau.de/publications/papers/2011/Preclik_Multibody_Ext_Abstr.pdf - """ - - # Dev note: - # Most of the cylinder-cylinder contact SHOULD be implemented - # as given in this `paper `, - # but the elastica-cpp kernels are implemented. - # This is maybe to speed-up the kernel, but it's - # potentially dangerous as it does not deal with "end" conditions - # correctly. - - def __init__( - self, - k: np.floating, - nu: np.floating, - velocity_damping_coefficient: np.floating = 0, - friction_coefficient: np.floating = 0, - ) -> None: - """ - - Parameters - ---------- - k : float - Contact spring constant. - nu : float - Contact damping constant. - velocity_damping_coefficient : float - Velocity damping coefficient between rigid-body and rod contact is used to apply friction force in the - slip direction. - friction_coefficient : float - For Coulombic friction coefficient for rigid-body and rod contact. - """ - super().__init__(k, nu) - self.velocity_damping_coefficient = velocity_damping_coefficient - self.friction_coefficient = friction_coefficient - log = logging.getLogger(self.__class__.__name__) - log.warning( - # Remove warning and add error if ExternalContact is used in v0.3.3 - # Remove the option to use ExternalContact, beyond v0.3.3 - "The option to use the ExternalContact joint for the rod-rod and rod-cylinder contact is now deprecated.\n" - "Instead, for rod-rod contact or rod-cylinder contact,use RodRodContact or RodCylinderContact from the add-on Contact mixin class.\n" - "For reference see the classes elastica.contact_forces.RodRodContact() and elastica.contact_forces.RodCylinderContact().\n" - "For usage check examples/RigidbodyCases/RodRigidBodyContact/rod_cylinder_contact.py and examples/RodContactCase/RodRodContact/rod_rod_contact_parallel_validation.py.\n" - " The option to use the ExternalContact joint for the rod-rod and rod-cylinder will be removed in the future (v0.3.3).\n" - ) - - def apply_forces( - self, - rod_one: SystemType, - index_one: ConnectionIndex, - rod_two: SystemType, - index_two: ConnectionIndex, - ) -> None: - # del index_one, index_two - from elastica.contact_utils import ( - _prune_using_aabbs_rod_cylinder, - _prune_using_aabbs_rod_rod, - ) - from elastica._contact_functions import ( - _calculate_contact_forces_rod_cylinder, - _calculate_contact_forces_rod_rod, - ) - - # TODO: raise error during the initialization if rod one is rigid body. - - # If rod two has one element, then it is rigid body. - if rod_two.n_elems == 1: - cylinder_two = rod_two - # First, check for a global AABB bounding box, and see whether that - # intersects - if _prune_using_aabbs_rod_cylinder( - rod_one.position_collection, - rod_one.radius, - rod_one.lengths, - cylinder_two.position_collection, - cylinder_two.director_collection, - cylinder_two.radius[0], - cylinder_two.length[0], - ): - return - - x_cyl = ( - cylinder_two.position_collection[..., 0] - - 0.5 * cylinder_two.length * cylinder_two.director_collection[2, :, 0] - ) - - rod_element_position = 0.5 * ( - rod_one.position_collection[..., 1:] - + rod_one.position_collection[..., :-1] - ) - _calculate_contact_forces_rod_cylinder( - rod_element_position, - rod_one.lengths * rod_one.tangents, - cylinder_two.position_collection[..., 0], - x_cyl, - cylinder_two.length * cylinder_two.director_collection[2, :, 0], - rod_one.radius + cylinder_two.radius, - rod_one.lengths + cylinder_two.length, - rod_one.internal_forces, - rod_one.external_forces, - cylinder_two.external_forces, - cylinder_two.external_torques, - cylinder_two.director_collection[:, :, 0], - rod_one.velocity_collection, - cylinder_two.velocity_collection, - self.k, - self.nu, - self.velocity_damping_coefficient, - self.friction_coefficient, - ) - - else: - # First, check for a global AABB bounding box, and see whether that - # intersects - - if _prune_using_aabbs_rod_rod( - rod_one.position_collection, - rod_one.radius, - rod_one.lengths, - rod_two.position_collection, - rod_two.radius, - rod_two.lengths, - ): - return - - _calculate_contact_forces_rod_rod( - rod_one.position_collection[ - ..., :-1 - ], # Discount last node, we want element start position - rod_one.radius, - rod_one.lengths, - rod_one.tangents, - rod_one.velocity_collection, - rod_one.internal_forces, - rod_one.external_forces, - rod_two.position_collection[ - ..., :-1 - ], # Discount last node, we want element start position - rod_two.radius, - rod_two.lengths, - rod_two.tangents, - rod_two.velocity_collection, - rod_two.internal_forces, - rod_two.external_forces, - self.k, - self.nu, - ) - - -class SelfContact(FreeJoint): - """ - This class is modeling self contact of rod. - - """ - - def __init__(self, k: np.floating, nu: np.floating) -> None: - super().__init__(k, nu) - log = logging.getLogger(self.__class__.__name__) - log.warning( - # Remove warning and add error if SelfContact is used in v0.3.3 - # Remove the option to use SelfContact, beyond v0.3.3 - "The option to use the SelfContact joint for the rod self contact is now deprecated.\n" - "Instead, for rod self contact use RodSelfContact from the add-on Contact mixin class.\n" - "For reference see the class elastica.contact_forces.RodSelfContact(), and for usage check examples/RodContactCase/RodSelfContact/solenoids.py.\n" - "The option to use the SelfContact joint for the rod self contact will be removed in the future (v0.3.3).\n" - ) - - def apply_forces( - self, - rod_one: SystemType, - index_one: ConnectionIndex, - rod_two: SystemType, - index_two: ConnectionIndex, - ) -> None: - # del index_one, index_two - from elastica._contact_functions import ( - _calculate_contact_forces_self_rod, - ) - - _calculate_contact_forces_self_rod( - rod_one.position_collection[ - ..., :-1 - ], # Discount last node, we want element start position - rod_one.radius, - rod_one.lengths, - rod_one.tangents, - rod_one.velocity_collection, - rod_one.external_forces, - self.k, - self.nu, - ) From 0b38d364769bdc368715cd51421a9a163283469c Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 16 Jun 2024 10:27:02 +0900 Subject: [PATCH 090/134] refactor: mv memory related methods into separate file --- elastica/systems/__init__.py | 90 ++---------------------------------- elastica/systems/memory.py | 84 +++++++++++++++++++++++++++++++++ 2 files changed, 87 insertions(+), 87 deletions(-) create mode 100644 elastica/systems/memory.py diff --git a/elastica/systems/__init__.py b/elastica/systems/__init__.py index 246629c39..936a3218a 100644 --- a/elastica/systems/__init__.py +++ b/elastica/systems/__init__.py @@ -1,11 +1,8 @@ -from typing import Iterator, TypeVar, Generic, Type -from elastica.timestepper.protocol import ExplicitStepperProtocol -from elastica.typing import SystemCollectionType +from typing import Type -from copy import copy - -def is_system_a_collection(system: object) -> bool: +# TODO: Be more specific Type +def is_system_a_collection(system: Type) -> bool: # Check if system is a "collection" of smaller systems # by checking for the [] method """ @@ -43,84 +40,3 @@ def is_system_a_collection(system: object) -> bool: return issubclass(system.__class__, BaseSystemCollection) or callable( __sys_get_item ) - - -# FIXME: Move memory related functions to separate module or as part of the timestepper - - -# TODO: Use MemoryProtocol -def make_memory_for_explicit_stepper( - stepper: ExplicitStepperProtocol, system: SystemCollectionType -) -> "MemoryCollection": - # TODO Automated logic (class creation, memory management logic) agnostic of stepper details (RK, AB etc.) - - from elastica.timestepper.explicit_steppers import ( - RungeKutta4, - EulerForward, - ) - - # is_this_system_a_collection = is_system_a_collection(system) - - memory_cls: Type - if RungeKutta4 in stepper.__class__.mro(): - # Bad way of doing it, introduces tight coupling - # this should rather be taken from the class itself - class MemoryRungeKutta4: - def __init__(self) -> None: - self.initial_state = None - self.k_1 = None - self.k_2 = None - self.k_3 = None - self.k_4 = None - - memory_cls = MemoryRungeKutta4 - elif EulerForward in stepper.__class__.mro(): - - class MemoryEulerForward: - def __init__(self) -> None: - self.initial_state = None - self.k = None - - memory_cls = MemoryEulerForward - else: - raise NotImplementedError("Making memory for other types not supported") - - return MemoryCollection(memory_cls(), len(system)) - - -M = TypeVar("M", bound="MemoryCollection") - - -class MemoryCollection(Generic[M]): - """Slots of memories for timestepper in a cohesive unit. - - A `MemoryCollection` object is meant to be used in conjunction - with a `SystemCollection`, where each independent `System` to - be integrated has its own `Memory`. - - Example - ------- - - A RK4 integrator needs to store k_1, k_2, k_3, k_4 (intermediate - results from four stages) for each `System`. The restriction for - having a memory slot arises because the `Systems` are usually - not independent of one another and may need communication after - every stage. - """ - - def __init__(self, memory: M, n_memory_slots: int): - super(MemoryCollection, self).__init__() - - self.__memories: list[M] = [] - for _ in range(n_memory_slots - 1): - self.__memories.append(copy(memory)) - self.__memories.append(memory) - - def __getitem__(self, idx: int) -> M: - return self.__memories[idx] - - def __len__(self) -> int: - return len(self.__memories) - - def __iter__(self) -> Iterator[M]: - return self.__memories.__iter__() diff --git a/elastica/systems/memory.py b/elastica/systems/memory.py new file mode 100644 index 000000000..c669be9b8 --- /dev/null +++ b/elastica/systems/memory.py @@ -0,0 +1,84 @@ +from typing import Iterator, TypeVar, Generic, Type +from elastica.timestepper.protocol import ExplicitStepperProtocol +from elastica.typing import SystemCollectionType + +from copy import copy + + +# FIXME: Move memory related functions to separate module or as part of the timestepper +# TODO: Use MemoryProtocol +def make_memory_for_explicit_stepper( + stepper: ExplicitStepperProtocol, system: SystemCollectionType +) -> "MemoryCollection": + # TODO Automated logic (class creation, memory management logic) agnostic of stepper details (RK, AB etc.) + + from elastica.timestepper.explicit_steppers import ( + RungeKutta4, + EulerForward, + ) + + # is_this_system_a_collection = is_system_a_collection(system) + + memory_cls: Type + if RungeKutta4 in stepper.__class__.mro(): + # Bad way of doing it, introduces tight coupling + # this should rather be taken from the class itself + class MemoryRungeKutta4: + def __init__(self) -> None: + self.initial_state = None + self.k_1 = None + self.k_2 = None + self.k_3 = None + self.k_4 = None + + memory_cls = MemoryRungeKutta4 + elif EulerForward in stepper.__class__.mro(): + + class MemoryEulerForward: + def __init__(self) -> None: + self.initial_state = None + self.k = None + + memory_cls = MemoryEulerForward + else: + raise NotImplementedError("Making memory for other types not supported") + + return MemoryCollection(memory_cls(), len(system)) + + +M = TypeVar("M", bound="MemoryCollection") + + +class MemoryCollection(Generic[M]): + """Slots of memories for timestepper in a cohesive unit. + + A `MemoryCollection` object is meant to be used in conjunction + with a `SystemCollection`, where each independent `System` to + be integrated has its own `Memory`. + + Example + ------- + + A RK4 integrator needs to store k_1, k_2, k_3, k_4 (intermediate + results from four stages) for each `System`. The restriction for + having a memory slot arises because the `Systems` are usually + not independent of one another and may need communication after + every stage. + """ + + def __init__(self, memory: M, n_memory_slots: int): + super(MemoryCollection, self).__init__() + + self.__memories: list[M] = [] + for _ in range(n_memory_slots - 1): + self.__memories.append(copy(memory)) + self.__memories.append(memory) + + def __getitem__(self, idx: int) -> M: + return self.__memories[idx] + + def __len__(self) -> int: + return len(self.__memories) + + def __iter__(self) -> Iterator[M]: + return self.__memories.__iter__() From 8f5af01604f87fa9c64ed16e4204cafc38b1d9e5 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 16 Jun 2024 12:12:16 -0500 Subject: [PATCH 091/134] resolve circular import --- elastica/dissipation.py | 10 +++++---- elastica/external_forces.py | 2 +- elastica/interaction.py | 2 +- elastica/joint.py | 33 ++++++++++++++--------------- tests/test_math/test_timestepper.py | 2 +- 5 files changed, 25 insertions(+), 24 deletions(-) diff --git a/elastica/dissipation.py b/elastica/dissipation.py index 95453beaf..a9733ece3 100644 --- a/elastica/dissipation.py +++ b/elastica/dissipation.py @@ -29,7 +29,7 @@ class DamperBase(ABC): """ - _system: RodType | RigidBodyType + _system: "RodType | RigidBodyType" # TODO typing can be made better def __init__(self, *args: Any, **kwargs: Any) -> None: @@ -43,7 +43,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: ) @property - def system(self) -> RodType | RigidBodyType: + def system(self) -> "RodType | RigidBodyType": """ get system (rod or rigid body) reference @@ -55,7 +55,9 @@ def system(self) -> RodType | RigidBodyType: return self._system @abstractmethod - def dampen_rates(self, system: RodType | RigidBodyType, time: np.floating) -> None: + def dampen_rates( + self, system: "RodType | RigidBodyType", time: np.floating + ) -> None: # TODO: In the future, we can remove rod and use self.system """ Dampen rates (velocity and/or omega) of a rod object. @@ -148,7 +150,7 @@ def __init__( * np.diagonal(self._system.inv_mass_second_moment_of_inertia).T ) - def dampen_rates(self, rod: RodType | RigidBodyType, time: np.floating) -> None: + def dampen_rates(self, rod: "RodType | RigidBodyType", time: np.floating) -> None: rod.velocity_collection[:] = ( rod.velocity_collection * self.translational_damping_coefficient ) diff --git a/elastica/external_forces.py b/elastica/external_forces.py index f3d3655d7..1b4124be0 100644 --- a/elastica/external_forces.py +++ b/elastica/external_forces.py @@ -91,7 +91,7 @@ def __init__( self.acc_gravity = acc_gravity def apply_forces( - self, system: RodType | RigidBodyType, time: np.floating = np.float64(0.0) + self, system: "RodType | RigidBodyType", time: np.floating = np.float64(0.0) ) -> None: self.compute_gravity_forces( self.acc_gravity, system.mass, system.external_forces diff --git a/elastica/interaction.py b/elastica/interaction.py index d5ec64c33..c13c34f7d 100644 --- a/elastica/interaction.py +++ b/elastica/interaction.py @@ -248,7 +248,7 @@ def __init__( # kinetic and static friction should separate functions # for now putting them together to figure out common variables def apply_forces( - self, system: RodType | RigidBodyType, time: np.floating = np.float64(0) + self, system: "RodType | RigidBodyType", time: np.floating = np.float64(0) ) -> None: """ Call numba implementation to apply friction forces diff --git a/elastica/joint.py b/elastica/joint.py index 6df1e47ea..806bd40c1 100644 --- a/elastica/joint.py +++ b/elastica/joint.py @@ -4,10 +4,9 @@ from typing import Optional from elastica._rotations import _inv_rotate -from elastica.typing import SystemType, RodType, ConnectionIndex -import numpy as np +from elastica.typing import SystemType, RodType, ConnectionIndex, RigidBodyType -pass +import numpy as np from numpy.typing import NDArray @@ -49,9 +48,9 @@ def __init__(self, k: np.floating, nu: np.floating) -> None: def apply_forces( self, - system_one: RodType | RigidBodyType, + system_one: "RodType | RigidBodyType", index_one: ConnectionIndex, - system_two: RodType | RigidBodyType, + system_two: "RodType | RigidBodyType", index_two: ConnectionIndex, ) -> None: """ @@ -92,9 +91,9 @@ def apply_forces( def apply_torques( self, - system_one: RodType | RigidBodyType, + system_one: "RodType | RigidBodyType", index_one: ConnectionIndex, - system_two: RodType | RigidBodyType, + system_two: "RodType | RigidBodyType", index_two: ConnectionIndex, ) -> None: """ @@ -173,18 +172,18 @@ def __init__( # Apply force is same as free joint def apply_forces( self, - system_one: RodType | RigidBodyType, + system_one: "RodType | RigidBodyType", index_one: ConnectionIndex, - system_two: RodType | RigidBodyType, + system_two: "RodType | RigidBodyType", index_two: ConnectionIndex, ) -> None: return super().apply_forces(system_one, index_one, system_two, index_two) def apply_torques( self, - system_one: RodType | RigidBodyType, + system_one: "RodType | RigidBodyType", index_one: ConnectionIndex, - system_two: RodType | RigidBodyType, + system_two: "RodType | RigidBodyType", index_two: ConnectionIndex, ) -> None: # current tangent direction of the `index_two` element of system two @@ -280,18 +279,18 @@ def __init__( # Apply force is same as free joint def apply_forces( self, - system_one: RodType | RigidBodyType, + system_one: "RodType | RigidBodyType", index_one: ConnectionIndex, - system_two: RodType | RigidBodyType, + system_two: "RodType | RigidBodyType", index_two: ConnectionIndex, ) -> None: return super().apply_forces(system_one, index_one, system_two, index_two) def apply_torques( self, - system_one: RodType | RigidBodyType, + system_one: "RodType | RigidBodyType", index_one: ConnectionIndex, - system_two: RodType | RigidBodyType, + system_two: "RodType | RigidBodyType", index_two: ConnectionIndex, ) -> None: # collect directors of systems one and two @@ -337,9 +336,9 @@ def apply_torques( def get_relative_rotation_two_systems( - system_one: RodType | RigidBodyType, + system_one: "RodType | RigidBodyType", index_one: ConnectionIndex, - system_two: RodType | RigidBodyType, + system_two: "RodType | RigidBodyType", index_two: ConnectionIndex, ) -> NDArray[np.floating]: """ diff --git a/tests/test_math/test_timestepper.py b/tests/test_math/test_timestepper.py index bf4b2fa2c..2afa6de65 100644 --- a/tests/test_math/test_timestepper.py +++ b/tests/test_math/test_timestepper.py @@ -245,7 +245,7 @@ def test_explicit_steppers(self, explicit_stepper): # Before stepping, let's extend the interface of the stepper # while providing memory slots - from elastica.systems import make_memory_for_explicit_stepper + from elastica.systems.memory import make_memory_for_explicit_stepper memory_collection = make_memory_for_explicit_stepper(stepper, collective_system) from elastica.timestepper import extend_stepper_interface From 4df6150dd96107457be134c9b33a8a2baa2d5630 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 16 Jun 2024 12:12:40 -0500 Subject: [PATCH 092/134] test: remove unittest for ExternalContact and SelfContact --- tests/test_joint.py | 327 -------------------------------------------- 1 file changed, 327 deletions(-) diff --git a/tests/test_joint.py b/tests/test_joint.py index 40cbd7c27..6d5165914 100644 --- a/tests/test_joint.py +++ b/tests/test_joint.py @@ -13,7 +13,6 @@ import numpy as np import pytest from scipy.spatial.transform import Rotation -from elastica.joint import ExternalContact, SelfContact # TODO: change tests and made them independent of rod, at least assigin hardcoded values for forces and torques @@ -412,329 +411,3 @@ def mock_rigid_body_init(self): self.external_forces = np.array([[0.0], [0.0], [0.0]]) self.external_torques = np.array([[0.0], [0.0], [0.0]]) self.velocity_collection = np.array([[0.0], [0.0], [0.0]]) - - -MockRod = type("MockRod", (RodBase,), {"__init__": mock_rod_init}) - -MockRigidBody = type( - "MockRigidBody", (RigidBodyBase,), {"__init__": mock_rigid_body_init} -) - - -class TestExternalContact: - def test_external_contact_rod_rigid_body_with_collision_with_k_without_nu_and_friction( - self, - ): - - "Testing External Contact wrapper with Collision with analytical verified values" - - mock_rod = MockRod() - mock_rigid_body = MockRigidBody() - ext_contact = ExternalContact(k=1.0, nu=0.0) - ext_contact.apply_forces(mock_rod, 0, mock_rigid_body, 1) - - """Details and reasoning about the values are given in 'test_contact_specific_functions.py/test_claculate_contact_forces_rod_rigid_body()'""" - assert_allclose( - mock_rod.external_forces, - np.array([[0.166666, 0.333333, 0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]), - atol=1e-6, - ) - - assert_allclose( - mock_rigid_body.external_forces, np.array([[-0.5], [0.0], [0.0]]), atol=1e-6 - ) - - assert_allclose( - mock_rigid_body.external_torques, np.array([[0.0], [0.0], [0.0]]), atol=1e-6 - ) - - def test_external_contact_rod_rigid_body_with_collision_with_nu_without_k_and_friction( - self, - ): - - "Testing External Contact wrapper with Collision with analytical verified values" - - mock_rod = MockRod() - "Moving rod towards the cylinder with a velocity of -1 in x-axis" - mock_rod.velocity_collection = np.array([[-1, 0, 0], [-1, 0, 0], [-1, 0, 0]]) - mock_rigid_body = MockRigidBody() - "Moving cylinder towards the rod with a velocity of 1 in x-axis" - mock_rigid_body.velocity_collection = np.array([[1], [0], [0]]) - ext_contact = ExternalContact(k=0.0, nu=1.0) - ext_contact.apply_forces(mock_rod, 0, mock_rigid_body, 1) - - """Details and reasoning about the values are given in 'test_contact_specific_functions.py/test_claculate_contact_forces_rod_rigid_body()'""" - assert_allclose( - mock_rod.external_forces, - np.array([[0.5, 1, 0], [0, 0, 0], [0, 0, 0]]), - atol=1e-6, - ) - - assert_allclose( - mock_rigid_body.external_forces, np.array([[-1.5], [0], [0]]), atol=1e-6 - ) - - assert_allclose( - mock_rigid_body.external_torques, np.array([[0.0], [0.0], [0.0]]), atol=1e-6 - ) - - def test_external_contact_rod_rigid_body_with_collision_with_k_and_nu_without_friction( - self, - ): - - "Testing External Contact wrapper with Collision with analytical verified values" - - mock_rod = MockRod() - "Moving rod towards the cylinder with a velocity of -1 in x-axis" - mock_rod.velocity_collection = np.array([[-1, 0, 0], [-1, 0, 0], [-1, 0, 0]]) - mock_rigid_body = MockRigidBody() - "Moving cylinder towards the rod with a velocity of 1 in x-axis" - mock_rigid_body.velocity_collection = np.array([[1], [0], [0]]) - ext_contact = ExternalContact(k=1.0, nu=1.0) - ext_contact.apply_forces(mock_rod, 0, mock_rigid_body, 1) - - """Details and reasoning about the values are given in 'test_contact_specific_functions.py/test_claculate_contact_forces_rod_rigid_body()'""" - assert_allclose( - mock_rod.external_forces, - np.array([[0.666666, 1.333333, 0], [0, 0, 0], [0, 0, 0]]), - atol=1e-6, - ) - - assert_allclose( - mock_rigid_body.external_forces, np.array([[-2], [0], [0]]), atol=1e-6 - ) - - assert_allclose( - mock_rigid_body.external_torques, np.array([[0.0], [0.0], [0.0]]), atol=1e-6 - ) - - def test_external_contact_rod_rigid_body_with_collision_with_k_and_nu_and_friction( - self, - ): - - "Testing External Contact wrapper with Collision with analytical verified values" - - mock_rod = MockRod() - "Moving rod towards the cylinder with a velocity of -1 in x-axis" - mock_rod.velocity_collection = np.array([[-1, 0, 0], [-1, 0, 0], [-1, 0, 0]]) - mock_rigid_body = MockRigidBody() - "Moving cylinder towards the rod with a velocity of 1 in x-axis" - mock_rigid_body.velocity_collection = np.array([[1], [0], [0]]) - ext_contact = ExternalContact( - k=1.0, nu=1.0, velocity_damping_coefficient=0.1, friction_coefficient=0.1 - ) - ext_contact.apply_forces(mock_rod, 0, mock_rigid_body, 1) - - """Details and reasoning about the values are given in 'test_contact_specific_functions.py/test_claculate_contact_forces_rod_rigid_body()'""" - assert_allclose( - mock_rod.external_forces, - np.array( - [ - [0.666666, 1.333333, 0], - [0.033333, 0.066666, 0], - [0.033333, 0.066666, 0], - ] - ), - atol=1e-6, - ) - - assert_allclose( - mock_rigid_body.external_forces, np.array([[-2], [-0.1], [-0.1]]), atol=1e-6 - ) - - assert_allclose( - mock_rigid_body.external_torques, np.array([[0.0], [0.0], [0.0]]), atol=1e-6 - ) - - def test_external_contact_rod_rigid_body_without_collision(self): - - "Testing External Contact wrapper without Collision with analytical verified values" - - mock_rod = MockRod() - mock_rigid_body = MockRigidBody() - ext_contact = ExternalContact(k=1.0, nu=0.5) - - """Setting rigid body position such that there is no collision""" - mock_rigid_body.position_collection = np.array([[400], [500], [600]]) - mock_rod_external_forces_before_execution = mock_rod.external_forces.copy() - mock_rigid_body_external_forces_before_execution = ( - mock_rigid_body.external_forces.copy() - ) - mock_rigid_body_external_torques_before_execution = ( - mock_rigid_body.external_torques.copy() - ) - ext_contact.apply_forces(mock_rod, 0, mock_rigid_body, 1) - - assert_allclose( - mock_rod.external_forces, mock_rod_external_forces_before_execution - ) - assert_allclose( - mock_rigid_body.external_forces, - mock_rigid_body_external_forces_before_execution, - ) - assert_allclose( - mock_rigid_body.external_torques, - mock_rigid_body_external_torques_before_execution, - ) - - def test_external_contact_with_two_rods_with_collision_with_k_without_nu(self): - - "Testing External Contact wrapper with two rods with analytical verified values" - "Test values have been copied from 'test_contact_specific_functions.py/test_calculate_contact_forces_rod_rod()'" - - mock_rod_one = MockRod() - mock_rod_two = MockRod() - mock_rod_two.position_collection = np.array([[4, 5, 6], [0, 0, 0], [0, 0, 0]]) - ext_contact = ExternalContact(k=1.0, nu=0.0) - ext_contact.apply_forces(mock_rod_one, 0, mock_rod_two, 0) - - assert_allclose( - mock_rod_one.external_forces, - np.array([[0, -0.666666, -0.333333], [0, 0, 0], [0, 0, 0]]), - atol=1e-6, - ) - assert_allclose( - mock_rod_two.external_forces, - np.array([[0.333333, 0.666666, 0], [0, 0, 0], [0, 0, 0]]), - atol=1e-6, - ) - - def test_external_contact_with_two_rods_with_collision_without_k_with_nu(self): - - "Testing External Contact wrapper with two rods with analytical verified values" - "Test values have been copied from 'test_contact_specific_functions.py/test_calculate_contact_forces_rod_rod()'" - - mock_rod_one = MockRod() - mock_rod_two = MockRod() - - """Moving the rods towards each other with a velocity of 1 along the x-axis.""" - mock_rod_one.velocity_collection = np.array([[1, 0, 0], [1, 0, 0], [1, 0, 0]]) - mock_rod_two.velocity_collection = np.array( - [[-1, 0, 0], [-1, 0, 0], [-1, 0, 0]] - ) - mock_rod_two.position_collection = np.array([[4, 5, 6], [0, 0, 0], [0, 0, 0]]) - ext_contact = ExternalContact(k=0.0, nu=1.0) - ext_contact.apply_forces(mock_rod_one, 0, mock_rod_two, 0) - - assert_allclose( - mock_rod_one.external_forces, - np.array( - [[0, -0.333333, -0.166666], [0, 0, 0], [0, 0, 0]], - ), - atol=1e-6, - ) - assert_allclose( - mock_rod_two.external_forces, - np.array([[0.166666, 0.333333, 0], [0, 0, 0], [0, 0, 0]]), - atol=1e-6, - ) - - def test_external_contact_with_two_rods_with_collision_with_k_and_nu(self): - - "Testing External Contact wrapper with two rods with analytical verified values" - "Test values have been copied from 'test_contact_specific_functions.py/test_calculate_contact_forces_rod_rod()'" - - mock_rod_one = MockRod() - mock_rod_two = MockRod() - - """Moving the rods towards each other with a velocity of 1 along the x-axis.""" - mock_rod_one.velocity_collection = np.array([[1, 0, 0], [1, 0, 0], [1, 0, 0]]) - mock_rod_two.velocity_collection = np.array( - [[-1, 0, 0], [-1, 0, 0], [-1, 0, 0]] - ) - mock_rod_two.position_collection = np.array([[4, 5, 6], [0, 0, 0], [0, 0, 0]]) - ext_contact = ExternalContact(k=1.0, nu=1.0) - ext_contact.apply_forces(mock_rod_one, 0, mock_rod_two, 0) - - assert_allclose( - mock_rod_one.external_forces, - np.array( - [[0, -1, -0.5], [0, 0, 0], [0, 0, 0]], - ), - atol=1e-6, - ) - assert_allclose( - mock_rod_two.external_forces, - np.array([[0.5, 1, 0], [0, 0, 0], [0, 0, 0]]), - atol=1e-6, - ) - - def test_external_contact_with_two_rods_without_collision(self): - - "Testing External Contact wrapper with two rods with analytical verified values" - - mock_rod_one = MockRod() - mock_rod_two = MockRod() - - "Setting rod two position such that there is no collision" - mock_rod_two.position_collection = np.array( - [[100, 101, 102], [0, 0, 0], [0, 0, 0]] - ) - ext_contact = ExternalContact(k=1.0, nu=1.0) - mock_rod_one_external_forces_before_execution = ( - mock_rod_one.external_forces.copy() - ) - mock_rod_two_external_forces_before_execution = ( - mock_rod_two.external_forces.copy() - ) - ext_contact.apply_forces(mock_rod_one, 0, mock_rod_two, 0) - - assert_allclose( - mock_rod_one.external_forces, mock_rod_one_external_forces_before_execution - ) - assert_allclose( - mock_rod_two.external_forces, mock_rod_two_external_forces_before_execution - ) - - -class TestSelfContact: - def test_self_contact_with_rod_self_collision(self): - - "Testing Self Contact wrapper rod self collision with analytical verified values" - - mock_rod = MockRod() - - "Test values have been copied from 'test_contact_specific_functions.py/test_calculate_contact_forces_self_rod()'" - mock_rod.n_elems = 3 - mock_rod.position_collection = np.array( - [[1, 4, 4, 1], [0, 0, 1, 1], [0, 0, 0, 0]] - ) - mock_rod.radius = np.array([1, 1, 1]) - mock_rod.lengths = np.array([3, 1, 3]) - mock_rod.tangents = np.array( - [[1.0, 0.0, -1.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]] - ) - mock_rod.velocity_collection = np.array( - [[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]] - ) - mock_rod.internal_forces = np.array( - [[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]] - ) - mock_rod.external_forces = np.array( - [[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]] - ) - sel_contact = SelfContact(k=1.0, nu=0.0) - sel_contact.apply_forces(mock_rod, 0, mock_rod, 0) - - assert_allclose( - mock_rod.external_forces, - np.array( - [[0, 0, 0, 0], [-0.333333, -0.666666, 0.666666, 0.333333], [0, 0, 0, 0]] - ), - atol=1e-6, - ) - - def test_self_contact_with_rod_no_self_collision(self): - - "Testing Self Contact wrapper rod no self collision with analytical verified values" - - mock_rod = MockRod() - - "the initially set rod does not have self collision" - mock_rod_external_forces_before_execution = mock_rod.external_forces.copy() - sel_contact = SelfContact(k=1.0, nu=1.0) - sel_contact.apply_forces(mock_rod, 0, mock_rod, 0) - - assert_allclose( - mock_rod.external_forces, mock_rod_external_forces_before_execution - ) From 0e1c9ab00f1a32f6ed65545da6b1231537294221 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 16 Jun 2024 12:29:11 -0500 Subject: [PATCH 093/134] typing: collision utils --- elastica/collision/AABBCollection.py | 47 +++++++++++++++++++--------- elastica/modules/protocol.py | 5 ++- elastica/utils.py | 4 +-- 3 files changed, 39 insertions(+), 17 deletions(-) diff --git a/elastica/collision/AABBCollection.py b/elastica/collision/AABBCollection.py index 22b716065..c83fccac0 100644 --- a/elastica/collision/AABBCollection.py +++ b/elastica/collision/AABBCollection.py @@ -1,17 +1,20 @@ """ Axis Aligned Bounding Boxes for coarse collision detection """ +from typing_extensions import Self + import numpy as np +from numpy.typing import NDArray from elastica.utils import MaxDimension class AABBCollection: def __init__( self, - elemental_position_collection, - dimension_collection, + elemental_position_collection: NDArray[np.floating], + dimension_collection: NDArray[np.floating], elements_per_aabb: int, - ): + ) -> None: """ Doesn't differentiate tangent direction from the rest : potentially harmful as maybe you don't need to expand to radius amount in tangential direction @@ -36,7 +39,9 @@ def __init__( self.update(elemental_position_collection, dimension_collection) @classmethod - def make_from_aabb(cls, aabb_collection, scale_factor=4): + def make_from_aabb( + cls, aabb_collection: list[AABBCollection], scale_factor: int = 4 + ) -> Self: # Make position collection and dimension collection arrays from aabb_collection # Wasted effort, but only once during construction n_aabb_from_lower_level = len(aabb_collection) @@ -59,7 +64,7 @@ def make_from_aabb(cls, aabb_collection, scale_factor=4): return cls(elemental_position_collection, dimension_collection, scale_factor) - def _update(self, aabb_collection): + def _update(self, aabb_collection: list[AABBCollection]) -> None: # Updates internal state from another aabb """ @@ -78,7 +83,11 @@ def _update(self, aabb_collection): temp = np.array([aabb.aabb[..., 1, 0] for aabb in aabb_collection]) self.aabb[..., 1, 0] = np.amax(temp, axis=0) - def update(self, elemental_position_collection, dimension_collection): + def update( + self, + elemental_position_collection: NDArray[np.floating], + dimension_collection: NDArray[np.floating], + ) -> None: # Initialize the boxes for i in range(self.n_aabb): start = i * self.elements_per_aabb @@ -91,7 +100,7 @@ def update(self, elemental_position_collection, dimension_collection): ) + np.amax(dimension_collection[..., start:stop], axis=1) -def find_nearest_integer_square_root(x: int): +def find_nearest_integer_square_root(x: int) -> int: from math import sqrt return round(sqrt(x)) @@ -101,8 +110,11 @@ class AABBHierarchy: """Simple hierarchy for handling cylinder collisions alone, meant for a rod""" def __init__( - self, position_collection, dimension_collection, avg_n_dofs_in_final_level - ): + self, + position_collection: NDArray[np.floating], + dimension_collection: NDArray[np.floating], + avg_n_dofs_in_final_level: int, + ) -> None: """ scaling is always set to 4, so that theres' 1 major AABBCollection, then scaling_factor smaller AABBs, then scaling factor even smaller AABBs (which cover the elements @@ -121,10 +133,10 @@ def __init__( ) # nearest power of 4 that is less than the number - n_levels_bound_below = np.int( + n_levels_bound_below = int( np.floor(0.5 * np.log2(potential_n_aabbs_in_final_level)) ) - n_levels_bound_above = np.int( + n_levels_bound_above = int( np.ceil(0.5 * np.log2(potential_n_aabbs_in_final_level)) ) # Check which is the closest and use that as the number of levels @@ -214,11 +226,15 @@ def __init__( # Add one for the middle level # self.aabb.append(AABBCollection(position_collection, dimension_collection, self.n_aabbs_in_first_level)) - def n_aabbs_at_level(self, i: int): + def n_aabbs_at_level(self, i: int) -> int: assert i < self.n_levels return 4 ** (i) - def update(self, position_collection, dimension_collection): + def update( + self, + position_collection: NDArray[np.floating], + dimension_collection: NDArray[np.floating], + ) -> None: # Update bottom level first, the first level entries n_aabbs_in_final_level = self.n_aabbs_at_level(self.n_levels - 1) stop = 0 @@ -261,5 +277,8 @@ def update(self, position_collection, dimension_collection): count_elapsed_n_aabbs += n_aabbs_in_next_level -def are_aabb_intersecting(first_aabb_collection, second_aabb_collection): +def are_aabb_intersecting( + first_aabb_collection: NDArray[np.floating], + second_aabb_collection: NDArray[np.floating], +) -> bool: return True diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index f27544650..95f7b23b6 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -1,6 +1,9 @@ from typing import Protocol, Generator, TypeVar, Any, Type from typing_extensions import Self # 3.11: from typing import Self + from abc import abstractmethod +from collections.abc import MutableSequence + from elastica.typing import ( SystemIdxType, OperatorType, @@ -30,7 +33,7 @@ def instantiate(self, *args: Any, **kwargs: Any) -> M: ... def id(self) -> Any: ... -class SystemCollectionProtocol(Protocol): +class SystemCollectionProtocol(MutableSequence, Protocol): _systems: list[SystemType] @property diff --git a/elastica/utils.py b/elastica/utils.py index d31cf3cdc..e8a5d3832 100644 --- a/elastica/utils.py +++ b/elastica/utils.py @@ -72,7 +72,7 @@ def value() -> Literal[3]: class Tolerance: @staticmethod - def atol() -> float: + def atol() -> np.float64: """ Static absolute tolerance method @@ -83,7 +83,7 @@ def atol() -> float: return finfo(float64).eps * 1e4 @staticmethod - def rtol() -> float: + def rtol() -> np.float64: """ Static relative tolerance method From 83ec4fe6470e09b83a23d68e740d45508758df69 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 16 Jun 2024 14:20:55 -0500 Subject: [PATCH 094/134] wip: typing symplectic stepper --- elastica/joint.py | 2 +- elastica/modules/base_system.py | 3 ++- elastica/modules/protocol.py | 3 ++- elastica/surface/plane.py | 2 +- elastica/systems/__init__.py | 8 +++---- elastica/timestepper/explicit_steppers.py | 2 +- elastica/timestepper/symplectic_steppers.py | 14 ++++++------- elastica/typing.py | 2 +- elastica/utils.py | 6 +++--- elastica/wrappers.py | 23 --------------------- 10 files changed, 22 insertions(+), 43 deletions(-) delete mode 100644 elastica/wrappers.py diff --git a/elastica/joint.py b/elastica/joint.py index 806bd40c1..812a754c9 100644 --- a/elastica/joint.py +++ b/elastica/joint.py @@ -239,7 +239,7 @@ def __init__( k: np.floating, nu: np.floating, kt: np.floating, - nut: np.floating = 0.0, + nut: np.floating = np.float64(0.0), rest_rotation_matrix: Optional[NDArray[np.floating]] = None, ) -> None: """ diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index 8e91fd1b3..74c1439bb 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -13,6 +13,7 @@ OperatorType, OperatorCallbackType, OperatorFinalizeType, + BlockType, ) import numpy as np @@ -68,7 +69,7 @@ def __init__(self) -> None: # List of systems to be integrated self._systems: list[SystemType] = [] - self._memory_blocks: list[SystemType] = [] + self._memory_blocks: list[BlockType] = [] # Flag Finalize: Finalizing twice will cause an error, # but the error message is very misleading diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index 95f7b23b6..58bedc2de 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -11,6 +11,7 @@ OperatorFinalizeType, SystemType, ConnectionIndex, + BlockType, ) from elastica.joint import FreeJoint from elastica.callback_functions import CallBackBaseClass @@ -59,7 +60,7 @@ def apply_callbacks(self, time: np.floating, current_step: int) -> None: ... @property def _feature_group_finalize(self) -> list[OperatorFinalizeType]: ... - def blocks(self) -> Generator[SystemType, None, None]: ... + def blocks(self) -> Generator[BlockType, None, None]: ... def _get_sys_idx_if_valid(self, sys_to_be_added: SystemType) -> SystemIdxType: ... diff --git a/elastica/surface/plane.py b/elastica/surface/plane.py index 82edd1a64..7f17537f0 100644 --- a/elastica/surface/plane.py +++ b/elastica/surface/plane.py @@ -23,7 +23,7 @@ def __init__(self, plane_origin: np.ndarray, plane_normal: np.ndarray): assert np.allclose( np.linalg.norm(plane_normal), 1, - atol=Tolerance.atol(), + atol=float(Tolerance.atol()), ), "plane normal is not a unit vector" self.normal = np.asarray(plane_normal).reshape(3) self.origin = np.asarray(plane_origin).reshape(3, 1) diff --git a/elastica/systems/__init__.py b/elastica/systems/__init__.py index 936a3218a..80ceb76ba 100644 --- a/elastica/systems/__init__.py +++ b/elastica/systems/__init__.py @@ -1,8 +1,10 @@ from typing import Type +from elastica.typing import SystemType, SystemCollectionType + # TODO: Be more specific Type -def is_system_a_collection(system: Type) -> bool: +def is_system_a_collection(system: SystemType | SystemCollectionType) -> bool: # Check if system is a "collection" of smaller systems # by checking for the [] method """ @@ -37,6 +39,4 @@ def is_system_a_collection(system: Type) -> bool: from elastica.modules import BaseSystemCollection __sys_get_item = getattr(system, "__getitem__", None) - return issubclass(system.__class__, BaseSystemCollection) or callable( - __sys_get_item - ) + return isinstance(system, BaseSystemCollection) or callable(__sys_get_item) diff --git a/elastica/timestepper/explicit_steppers.py b/elastica/timestepper/explicit_steppers.py index aeb4cab03..5a5877f67 100644 --- a/elastica/timestepper/explicit_steppers.py +++ b/elastica/timestepper/explicit_steppers.py @@ -188,7 +188,7 @@ def _first_update( time: np.floating, dt: np.floating, ) -> np.floating: - System.state += dt * System(time, dt) + System.state += dt * System(time, dt) # type: ignore[arg-type] return time + dt diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index 84a87ce20..11184d1df 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -95,7 +95,7 @@ def do_step( """ for kin_prefactor, kin_step, dyn_step in steps_and_prefactors[:-1]: - for system in SystemCollection._memory_blocks: + for system in SystemCollection.blocks(): kin_step(system, time, dt) time += kin_prefactor(dt) @@ -104,14 +104,14 @@ def do_step( SystemCollection.constrain_values(time) # We need internal forces and torques because they are used by interaction module. - for system in SystemCollection._memory_blocks: + for system in SystemCollection.blocks(): system.update_internal_forces_and_torques(time) # system.update_internal_forces_and_torques() # Add external forces, controls etc. SystemCollection.synchronize(time) - for system in SystemCollection._memory_blocks: + for system in SystemCollection.blocks(): dyn_step(system, time, dt) # Constrain only rates @@ -121,7 +121,7 @@ def do_step( last_kin_prefactor = steps_and_prefactors[-1][0] last_kin_step = steps_and_prefactors[-1][1] - for system in SystemCollection._memory_blocks: + for system in SystemCollection.blocks(): last_kin_step(system, time, dt) time += last_kin_prefactor(dt) SystemCollection.constrain_values(time) @@ -130,14 +130,14 @@ def do_step( SystemCollection.apply_callbacks(time, round(time / dt)) # Zero out the external forces and torques - for system in SystemCollection._memory_blocks: + for system in SystemCollection.blocks(): system.reset_external_forces_and_torques(time) return time def step_single_instance( self: SymplecticStepperProtocol, - System: SystemType, + System: SymplecticSystemProtocol, time: np.floating, dt: np.floating, ) -> np.floating: @@ -145,7 +145,7 @@ def step_single_instance( for kin_prefactor, kin_step, dyn_step in self.steps_and_prefactors[:-1]: kin_step(System, time, dt) time += kin_prefactor(dt) - System.update_internal_forces_and_torques(time) + System.compute_internal_forces_and_torques(time) dyn_step(System, time, dt) # Peel the last kinematic step and prefactor alone diff --git a/elastica/typing.py b/elastica/typing.py index 8e94bb3b4..1cdf79dcd 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -39,7 +39,7 @@ from .joint import FreeJoint -SystemType: TypeAlias = "SymplecticSystemProtocol" # | ExplicitSystemProtocol +SystemType: TypeAlias = "SystemProtocol" SystemIdxType: TypeAlias = int # ModuleObjectTypes: TypeAlias = ( diff --git a/elastica/utils.py b/elastica/utils.py index e8a5d3832..dde1e869f 100644 --- a/elastica/utils.py +++ b/elastica/utils.py @@ -178,8 +178,8 @@ def extend_instance(obj: Any, cls: Any) -> None: obj.__class__ = type(base_cls_name, (cls, base_cls), {}) -def _bspline( - t_coeff: NDArray, l_centerline: np.floating = 1.0 +def _bspline( # type: ignore[no-any-unimported] + t_coeff: NDArray, l_centerline: np.floating = np.float64(1.0) ) -> tuple[BSpline, NDArray, NDArray]: """Generates a bspline object that plots the spline interpolant for any vector x. Optionally takes in a centerline length, set to 1.0 by @@ -208,7 +208,7 @@ def _bspline( return __bspline_impl__(control_pts, t_coeff, degree) -def __bspline_impl__( +def __bspline_impl__( # type: ignore[no-any-unimported] x_pts: NDArray, t_c: NDArray, degree: int ) -> tuple[BSpline, NDArray, NDArray]: """""" diff --git a/elastica/wrappers.py b/elastica/wrappers.py deleted file mode 100644 index fa3990f80..000000000 --- a/elastica/wrappers.py +++ /dev/null @@ -1,23 +0,0 @@ -import warnings - -__all__ = [ - "BaseSystemCollection", - "Connections", - "Constraints", - "Forcing", - "CallBacks", - "Damping", -] - -from elastica.modules.base_system import BaseSystemCollection -from elastica.modules.connections import Connections -from elastica.modules.constraints import Constraints -from elastica.modules.forcing import Forcing -from elastica.modules.callbacks import CallBacks -from elastica.modules.damping import Damping - - -warnings.warn( - "elastica.wrappers is refactored to elastica.modules in version 0.3.0.", - DeprecationWarning, -) From 267bca228733511620283686f1417f65f089a1cd Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 16 Jun 2024 14:26:49 -0500 Subject: [PATCH 095/134] fix tests with refactoring symplectic system --- elastica/modules/protocol.py | 7 ++++--- elastica/systems/__init__.py | 3 +-- elastica/systems/analytical.py | 3 +++ tests/test_rigid_body/test_rigid_body_data_structures.py | 2 +- 4 files changed, 9 insertions(+), 6 deletions(-) diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index 58bedc2de..5ffdcfaaf 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -1,8 +1,9 @@ -from typing import Protocol, Generator, TypeVar, Any, Type +from typing import Protocol, Generator, TypeVar, Any, Type, Sized from typing_extensions import Self # 3.11: from typing import Self from abc import abstractmethod -from collections.abc import MutableSequence + +pass from elastica.typing import ( SystemIdxType, @@ -34,7 +35,7 @@ def instantiate(self, *args: Any, **kwargs: Any) -> M: ... def id(self) -> Any: ... -class SystemCollectionProtocol(MutableSequence, Protocol): +class SystemCollectionProtocol(Sized, Protocol): _systems: list[SystemType] @property diff --git a/elastica/systems/__init__.py b/elastica/systems/__init__.py index 80ceb76ba..b61b905ed 100644 --- a/elastica/systems/__init__.py +++ b/elastica/systems/__init__.py @@ -3,8 +3,7 @@ from elastica.typing import SystemType, SystemCollectionType -# TODO: Be more specific Type -def is_system_a_collection(system: SystemType | SystemCollectionType) -> bool: +def is_system_a_collection(system: "SystemType | SystemCollectionType") -> bool: # Check if system is a "collection" of smaller systems # by checking for the [] method """ diff --git a/elastica/systems/analytical.py b/elastica/systems/analytical.py index b03a917aa..fba1b4a1b 100644 --- a/elastica/systems/analytical.py +++ b/elastica/systems/analytical.py @@ -302,6 +302,9 @@ def __init__(self): self._memory_blocks = [] self.systems = self._memory_blocks + def blocks(self): + return self._memory_blocks + def __getitem__(self, idx): return self._memory_blocks[idx] diff --git a/tests/test_rigid_body/test_rigid_body_data_structures.py b/tests/test_rigid_body/test_rigid_body_data_structures.py index b482202b8..0ca7715a0 100644 --- a/tests/test_rigid_body/test_rigid_body_data_structures.py +++ b/tests/test_rigid_body/test_rigid_body_data_structures.py @@ -64,7 +64,7 @@ def __init__(self, start_position, start_director): # Givees position, director etc. super(SimpleSystemWithPositionsDirectors, self).__init__() - def _compute_internal_forces_and_torques(self, time): + def compute_internal_forces_and_torques(self, time): pass def update_accelerations(self, time): From e4d9a1340617f46dd490b8971aabbd08e4cc56cb Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 16 Jun 2024 14:57:43 -0500 Subject: [PATCH 096/134] remove deprecated functions in interaction --- elastica/experimental/interaction.py | 6 +- elastica/interaction.py | 84 +---------- elastica/memory_block/protocol.py | 5 +- elastica/modules/connections.py | 15 +- elastica/rigidbody/protocol.py | 18 +++ elastica/rod/protocol.py | 2 - elastica/typing.py | 8 +- tests/test_interaction.py | 206 ++++----------------------- 8 files changed, 63 insertions(+), 281 deletions(-) create mode 100644 elastica/rigidbody/protocol.py diff --git a/elastica/experimental/interaction.py b/elastica/experimental/interaction.py index 776f0da7e..5609eba38 100644 --- a/elastica/experimental/interaction.py +++ b/elastica/experimental/interaction.py @@ -3,8 +3,8 @@ import numpy as np from elastica.external_forces import NoForces +from elastica.contact_utils import _find_slipping_elements from elastica.interaction import ( - find_slipping_elements, apply_normal_force_numba_rigid_body, InteractionPlaneRigidBody, ) @@ -122,7 +122,7 @@ def anisotropic_friction_numba_rigid_body( + kinetic_mu_backward * (1 - velocity_sign_along_axial_direction) ) # Call slip function to check if elements slipping or not - slip_function_along_axial_direction = find_slipping_elements( + slip_function_along_axial_direction = _find_slipping_elements( velocity_along_axial_direction, slip_velocity_tol ) kinetic_friction_force_along_axial_direction = -( @@ -151,7 +151,7 @@ def anisotropic_friction_numba_rigid_body( + kinetic_mu_backward * (1 - velocity_sign_along_binormal_direction) ) # Call slip function to check if elements slipping or not - slip_function_along_binormal_direction = find_slipping_elements( + slip_function_along_binormal_direction = _find_slipping_elements( velocity_along_binormal_direction, slip_velocity_tol ) kinetic_friction_force_along_binormal_direction = -( diff --git a/elastica/interaction.py b/elastica/interaction.py index c13c34f7d..5e5339865 100644 --- a/elastica/interaction.py +++ b/elastica/interaction.py @@ -20,42 +20,6 @@ from elastica.typing import SystemType, RodType, RigidBodyType -def find_slipping_elements(velocity_slip: Any, velocity_threshold: Any) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._find_slipping_elements()\n" - "instead for finding slipping elements." - ) - - -def node_to_element_mass_or_force(input: Any) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._node_to_element_mass_or_force()\n" - "instead for converting the mass/forces on rod nodes to elements." - ) - - -def nodes_to_elements(input: Any) -> NoReturn: - # Remove the function beyond v0.4.0 - raise NotImplementedError( - "This function is removed in v0.3.1. Please use\n" - "elastica.interaction.node_to_element_mass_or_force()\n" - "instead for node-to-element interpolation of mass/forces." - ) - - -@njit(cache=True) # type: ignore -def elements_to_nodes_inplace( - vector_in_element_frame: Any, vector_in_node_frame: Any -) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._elements_to_nodes_inplace()\n" - "instead for updating nodal forces using the forces computed on elements." - ) - - # base class for interaction # only applies normal force no friction class InteractionPlane: @@ -144,33 +108,13 @@ def apply_normal_force( ) -def apply_normal_force_numba( - plane_origin: Any, - plane_normal: Any, - surface_tol: Any, - k: Any, - nu: Any, - radius: Any, - mass: Any, - position_collection: Any, - velocity_collection: Any, - internal_forces: Any, - external_forces: Any, -) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. For rod plane contact please use: \n" - "elastica._contact_functions._calculate_contact_forces_rod_plane() \n" - "For detail, refer to issue #113." - ) - - # class for anisotropic frictional plane # NOTE: friction coefficients are passed as arrays in the order # mu_forward : mu_backward : mu_sideways # head is at x[0] and forward means head to tail # same convention for kinetic and static # mu named as to which direction it opposes -class AnisotropicFrictionalPlane(NoForces, InteractionPlane): +class AnisotropicFrictionalPlane(InteractionPlane, NoForces): """ This anisotropic friction plane class is for computing anisotropic friction forces on rods. @@ -358,32 +302,6 @@ def sum_over_elements(input: NDArray[np.floating]) -> np.floating: return output -def node_to_element_position(node_position_collection: Any) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. For node-to-element_position() interpolation please use: \n" - "elastica.contact_utils._node_to_element_position() for rod position \n" - "For detail, refer to issue #113." - ) - - -def node_to_element_velocity(mass: Any, node_velocity_collection: Any) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. For node-to-element_velocity() interpolation please use: \n" - "elastica.contact_utils._node_to_element_velocity() for rod velocity. \n" - "For detail, refer to issue #113." - ) - - -def node_to_element_pos_or_vel(vector_in_node_frame: Any) -> NoReturn: - # Remove the function beyond v0.4.0 - raise NotImplementedError( - "This function is removed in v0.3.0. For node-to-element interpolation please use: \n" - "elastica.contact_utils._node_to_element_position() for rod position \n" - "elastica.contact_utils._node_to_element_velocity() for rod velocity. \n" - "For detail, refer to issue #80." - ) - - @njit(cache=True) # type: ignore def slender_body_forces( tangents: NDArray[np.floating], diff --git a/elastica/memory_block/protocol.py b/elastica/memory_block/protocol.py index 5cf002462..80d9c7a98 100644 --- a/elastica/memory_block/protocol.py +++ b/elastica/memory_block/protocol.py @@ -1,6 +1,7 @@ from typing import Protocol from elastica.rod.protocol import CosseratRodProtocol +from elastica.rigid_body.protocol import RigidBodyProtocol from elastica.systems.protocol import SymplecticSystemProtocol @@ -8,5 +9,5 @@ class BlockCosseratRodProtocol(CosseratRodProtocol, SymplecticSystemProtocol, Pr pass -# class BlockRigidBodyProtocol(RigidBodyProtocol, SymplecticSystemProtocol, Protocol): -# pass +class BlockRigidBodyProtocol(RigidBodyProtocol, SymplecticSystemProtocol, Protocol): + pass diff --git a/elastica/modules/connections.py b/elastica/modules/connections.py index a9ffb19d8..2b82bb7ce 100644 --- a/elastica/modules/connections.py +++ b/elastica/modules/connections.py @@ -10,8 +10,9 @@ from elastica.typing import ( SystemIdxType, OperatorFinalizeType, - SystemType, ConnectionIndex, + RodType, + RigidBodyType, ) import numpy as np import functools @@ -39,8 +40,8 @@ def __init__(self: SystemCollectionProtocol) -> None: def connect( self: SystemCollectionProtocol, - first_rod: SystemType, - second_rod: SystemType, + first_rod: "RodType | RigidBodyType", + second_rod: "RodType | RigidBodyType", first_connect_idx: ConnectionIndex = None, second_connect_idx: ConnectionIndex = None, ) -> ModuleProtocol: @@ -51,9 +52,9 @@ def connect( Parameters ---------- - first_rod : SystemType + first_rod : RodType | RigidBodyType Rod-like object - second_rod : SystemType + second_rod : RodType | RigidBodyType Rod-like object first_connect_idx : Optional[int] Index of first rod for joint. @@ -90,9 +91,9 @@ def _finalize_connections(self: SystemCollectionProtocol) -> None: def apply_forces_and_torques( time: np.floating, connect_instance: FreeJoint, - system_one: SystemType, + system_one: "RodType | RigidBodyType", first_connect_idx: ConnectionIndex, - system_two: SystemType, + system_two: "RodType | RigidBodyType", second_connect_idx: ConnectionIndex, ) -> None: connect_instance.apply_forces( diff --git a/elastica/rigidbody/protocol.py b/elastica/rigidbody/protocol.py new file mode 100644 index 000000000..abecdecb1 --- /dev/null +++ b/elastica/rigidbody/protocol.py @@ -0,0 +1,18 @@ +from typing import Protocol + +import numpy as np +from numpy.typing import NDArray + +from elastica.systems.protocol import SystemProtocol + + +class RigidBodyProtocol(SystemProtocol, Protocol): + + mass: NDArray[np.floating] + volume: NDArray[np.floating] + length: np.floating + tangents: NDArray[np.floating] + radius: np.floating + + mass_second_moment_of_inertia: NDArray[np.floating] + inv_mass_second_moment_of_inertia: NDArray[np.floating] diff --git a/elastica/rod/protocol.py b/elastica/rod/protocol.py index 672b5a33e..e1a6c0552 100644 --- a/elastica/rod/protocol.py +++ b/elastica/rod/protocol.py @@ -3,8 +3,6 @@ import numpy as np from numpy.typing import NDArray -pass - from elastica.systems.protocol import SystemProtocol diff --git a/elastica/typing.py b/elastica/typing.py index 1cdf79dcd..c23be3ee4 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -14,13 +14,17 @@ # Used for type hinting without circular imports # NEVER BACK-IMPORT ANY ELASTICA MODULES HERE from .rod.protocol import CosseratRodProtocol - from .rigidbody.rigid_body import RigidBodyBase + from .rigidbody.protocol import RigidBodyProtocol from .surface.surface_base import SurfaceBase from .modules.base_system import BaseSystemCollection from .modules.protocol import SystemCollectionProtocol from .rod.data_structures import _State as State - from .systems.protocol import SymplecticSystemProtocol, ExplicitSystemProtocol + from .systems.protocol import ( + SystemProtocol, + SymplecticSystemProtocol, + ExplicitSystemProtocol, + ) from .timestepper.protocol import ( StepperProtocol, SymplecticStepperProtocol, diff --git a/tests/test_interaction.py b/tests/test_interaction.py index 8d2b0e13c..4a71053f7 100644 --- a/tests/test_interaction.py +++ b/tests/test_interaction.py @@ -6,12 +6,8 @@ from elastica.utils import Tolerance, MaxDimension from elastica.interaction import ( InteractionPlane, - find_slipping_elements, AnisotropicFrictionalPlane, - node_to_element_mass_or_force, SlenderBodyTheory, - nodes_to_elements, - elements_to_nodes_inplace, apply_normal_force_numba_rigid_body, ) from elastica.contact_utils import ( @@ -322,53 +318,6 @@ def test_interaction_when_rod_is_under_plane_without_k_with_nu(self, n_elem, nu_ assert_allclose(correct_forces, rod.external_forces, atol=Tolerance.atol()) -class TestAuxiliaryFunctions: - @pytest.mark.parametrize("n_elem", [2, 3, 5, 10, 20]) - def test_linear_interpolation_slip_error_message(self, n_elem): - velocity_threshold = 1.0 - - # if slip velocity larger than threshold - velocity_slip = np.repeat( - np.array([0.0, 0.0, 2.0]).reshape(3, 1), n_elem, axis=1 - ) - error_message = ( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._find_slipping_elements()\n" - "instead for finding slipping elements." - ) - with pytest.raises(NotImplementedError) as error_info: - slip_function = find_slipping_elements(velocity_slip, velocity_threshold) - assert error_info.value.args[0] == error_message - - @pytest.mark.parametrize("n_elem", [2, 3, 5, 10, 20]) - def test_node_to_element_mass_or_force_error_message(self, n_elem): - random_vector = np.random.rand(3).reshape(3, 1) - input = np.repeat(random_vector, n_elem + 1, axis=1) - input[..., 0] *= 0.5 - input[..., -1] *= 0.5 - error_message = ( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._node_to_element_mass_or_force()\n" - "instead for converting the mass/forces on rod nodes to elements." - ) - with pytest.raises(NotImplementedError) as error_info: - output = node_to_element_mass_or_force(input) - assert error_info.value.args[0] == error_message - - @pytest.mark.parametrize("n_elem", [2, 10]) - def test_not_impl_error_for_nodes_to_elements(self, n_elem): - random_vector = np.random.rand(3).reshape(3, 1) - input = np.repeat(random_vector, n_elem + 1, axis=1) - error_message = ( - "This function is removed in v0.3.1. Please use\n" - "elastica.interaction.node_to_element_mass_or_force()\n" - "instead for node-to-element interpolation of mass/forces." - ) - with pytest.raises(NotImplementedError) as error_info: - nodes_to_elements(input) - assert error_info.value.args[0] == error_message - - class TestAnisotropicFriction: def initializer( self, @@ -739,140 +688,33 @@ def test_static_rolling_friction_total_torque_larger_than_static_friction_force( # Slender Body Theory Unit Tests +from elastica.interaction import ( + sum_over_elements, +) -try: - from elastica.interaction import ( - sum_over_elements, - node_to_element_position, - node_to_element_velocity, - node_to_element_pos_or_vel, - ) - # These functions are used in the case if Numba is available - class TestAuxiliaryFunctionsForSlenderBodyTheory: - @pytest.mark.parametrize("n_elem", [2, 3, 5, 10, 20]) - def test_sum_over_elements(self, n_elem): - """ - This function test sum over elements function with - respect to default python function .sum(). We write - this function because with numba we can get the sum - faster. - Parameters - ---------- - n_elem - - Returns - ------- - - """ - - input_variable = np.random.rand(n_elem) - correct_output = input_variable.sum() - output = sum_over_elements(input_variable) - assert_allclose(correct_output, output, atol=Tolerance.atol()) - - @pytest.mark.parametrize("n_elem", [2, 3, 5, 10, 20]) - def test_node_to_element_position_error_message(self, n_elem): - """ - This function tests node_to_element_position function. We are - converting node positions to element positions. Here also - we are using numba to speed up the process. - - Parameters - ---------- - n_elem - - Returns - ------- - - """ - random = np.random.rand() # Adding some random numbers - input_position = random * np.ones((3, n_elem + 1)) - correct_output = random * np.ones((3, n_elem)) - - error_message = ( - "This function is removed in v0.3.2. For node-to-element_position() interpolation please use: \n" - "elastica.contact_utils._node_to_element_position() for rod position \n" - "For detail, refer to issue #113." - ) - with pytest.raises(NotImplementedError) as error_info: - output = node_to_element_position(input_position) - assert error_info.value.args[0] == error_message - - @pytest.mark.parametrize("n_elem", [2, 3, 5, 10, 20]) - def test_node_to_element_velocity_error_message(self, n_elem): - """ - This function tests node_to_element_velocity function. We are - converting node velocities to element velocities. Here also - we are using numba to speed up the process. - - Parameters - ---------- - n_elem - - Returns - ------- - - """ - random = np.random.rand() # Adding some random numbers - input_velocity = random * np.ones((3, n_elem + 1)) - input_mass = 2.0 * random * np.ones(n_elem + 1) - correct_output = random * np.ones((3, n_elem)) - - error_message = ( - "This function is removed in v0.3.2. For node-to-element_velocity() interpolation please use: \n" - "elastica.contact_utils._node_to_element_velocity() for rod velocity. \n" - "For detail, refer to issue #113." - ) - with pytest.raises(NotImplementedError) as error_info: - output = node_to_element_velocity( - mass=input_mass, node_velocity_collection=input_velocity - ) - assert error_info.value.args[0] == error_message - - @pytest.mark.parametrize("n_elem", [2, 3, 5, 10, 20]) - def test_not_impl_error_for_node_to_element_pos_or_vel(self, n_elem): - random = np.random.rand() # Adding some random numbers - input_velocity = random * np.ones((3, n_elem + 1)) - error_message = ( - "This function is removed in v0.3.0. For node-to-element interpolation please use: \n" - "elastica.contact_utils._node_to_element_position() for rod position \n" - "elastica.contact_utils._node_to_element_velocity() for rod velocity. \n" - "For detail, refer to issue #80." - ) - with pytest.raises(NotImplementedError) as error_info: - node_to_element_pos_or_vel(input_velocity) - assert error_info.value.args[0] == error_message - - @pytest.mark.parametrize("n_elem", [2, 3, 5, 10, 20]) - def test_elements_to_nodes_inplace_error_message(self, n_elem): - """ - This function tests _elements_to_nodes_inplace. We are - converting node velocities to element velocities. Here also - we are using numba to speed up the process. - - Parameters - ---------- - n_elem - - Returns - ------- - - """ - random = np.random.rand() # Adding some random numbers - vector_in_element_frame = random * np.ones((3, n_elem)) - vector_in_node_frame = np.zeros((3, n_elem + 1)) - error_message = ( - "This function is removed in v0.3.2. Please use\n" - "elastica.contact_utils._elements_to_nodes_inplace()\n" - "instead for updating nodal forces using the forces computed on elements." - ) - with pytest.raises(NotImplementedError) as error_info: - elements_to_nodes_inplace(vector_in_element_frame, vector_in_node_frame) - assert error_info.value.args[0] == error_message +# These functions are used in the case if Numba is available +class TestAuxiliaryFunctionsForSlenderBodyTheory: + @pytest.mark.parametrize("n_elem", [2, 3, 5, 10, 20]) + def test_sum_over_elements(self, n_elem): + """ + This function test sum over elements function with + respect to default python function .sum(). We write + this function because with numba we can get the sum + faster. + Parameters + ---------- + n_elem + + Returns + ------- + + """ -except ImportError: - pass + input_variable = np.random.rand(n_elem) + correct_output = input_variable.sum() + output = sum_over_elements(input_variable) + assert_allclose(correct_output, output, atol=Tolerance.atol()) class TestSlenderBody: From 090ac1ea29ea6151c05244cd454a58954a735507 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 16 Jun 2024 15:35:51 -0500 Subject: [PATCH 097/134] remove deprecated anisotropic friction msg --- elastica/callback_functions.py | 1 + elastica/interaction.py | 32 -------------------------------- 2 files changed, 1 insertion(+), 32 deletions(-) diff --git a/elastica/callback_functions.py b/elastica/callback_functions.py index 88b15788d..4e885dbee 100644 --- a/elastica/callback_functions.py +++ b/elastica/callback_functions.py @@ -5,6 +5,7 @@ import os import sys import numpy as np +from numpy.typing import NDArray import logging diff --git a/elastica/interaction.py b/elastica/interaction.py index 5e5339865..c8f69eee5 100644 --- a/elastica/interaction.py +++ b/elastica/interaction.py @@ -229,38 +229,6 @@ def apply_forces( ) -def anisotropic_friction( - plane_origin: Any, - plane_normal: Any, - surface_tol: Any, - slip_velocity_tol: Any, - k: Any, - nu: Any, - kinetic_mu_forward: Any, - kinetic_mu_backward: Any, - kinetic_mu_sideways: Any, - static_mu_forward: Any, - static_mu_backward: Any, - static_mu_sideways: Any, - radius: Any, - mass: Any, - tangents: Any, - position_collection: Any, - director_collection: Any, - velocity_collection: Any, - omega_collection: Any, - internal_forces: Any, - external_forces: Any, - internal_torques: Any, - external_torques: Any, -) -> NoReturn: - raise NotImplementedError( - "This function is removed in v0.3.2. For anisotropic_friction please use: \n" - "elastica._contact_functions._calculate_contact_forces_rod_plane_with_anisotropic_friction() \n" - "For detail, refer to issue #113." - ) - - # Slender body module @njit(cache=True) # type: ignore def sum_over_elements(input: NDArray[np.floating]) -> np.floating: From f6738e97501ae954bca5a1a24b9520013e370c7d Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 16 Jun 2024 15:36:45 -0500 Subject: [PATCH 098/134] wip: fixing type disagreements --- elastica/callback_functions.py | 4 +++- elastica/dissipation.py | 14 ++++++-------- elastica/mesh/protocol.py | 10 ++++++++++ elastica/modules/protocol.py | 6 ++---- elastica/rigidbody/data_structures.py | 16 ++++------------ elastica/rigidbody/mesh_rigid_body.py | 21 ++++++++++++--------- elastica/typing.py | 5 ++++- elastica/utils.py | 8 ++++---- 8 files changed, 45 insertions(+), 39 deletions(-) create mode 100644 elastica/mesh/protocol.py diff --git a/elastica/callback_functions.py b/elastica/callback_functions.py index 4e885dbee..cc03be68b 100644 --- a/elastica/callback_functions.py +++ b/elastica/callback_functions.py @@ -178,7 +178,9 @@ def __init__( self.file_save_interval = file_save_interval # Data collector - self.buffer = defaultdict(list) + self.buffer: dict[str, list[NDArray[np.floating] | np.floating | int]] = ( + defaultdict(list) + ) self.buffer_size = 0 # Module diff --git a/elastica/dissipation.py b/elastica/dissipation.py index a9733ece3..7592c85b4 100644 --- a/elastica/dissipation.py +++ b/elastica/dissipation.py @@ -7,7 +7,7 @@ from abc import ABC, abstractmethod from typing import Any -from elastica.typing import RodType, RigidBodyType, SystemType +from elastica.typing import RodType, SystemType from numba import njit @@ -25,11 +25,11 @@ class DamperBase(ABC): Attributes ---------- - system : SystemType (RodBase or RigidBodyBase) + system : RodBase """ - _system: "RodType | RigidBodyType" + _system: RodType # TODO typing can be made better def __init__(self, *args: Any, **kwargs: Any) -> None: @@ -43,7 +43,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: ) @property - def system(self) -> "RodType | RigidBodyType": + def system(self) -> RodType: """ get system (rod or rigid body) reference @@ -55,9 +55,7 @@ def system(self) -> "RodType | RigidBodyType": return self._system @abstractmethod - def dampen_rates( - self, system: "RodType | RigidBodyType", time: np.floating - ) -> None: + def dampen_rates(self, system: RodType, time: np.floating) -> None: # TODO: In the future, we can remove rod and use self.system """ Dampen rates (velocity and/or omega) of a rod object. @@ -150,7 +148,7 @@ def __init__( * np.diagonal(self._system.inv_mass_second_moment_of_inertia).T ) - def dampen_rates(self, rod: "RodType | RigidBodyType", time: np.floating) -> None: + def dampen_rates(self, rod: RodType, time: np.floating) -> None: rod.velocity_collection[:] = ( rod.velocity_collection * self.translational_damping_coefficient ) diff --git a/elastica/mesh/protocol.py b/elastica/mesh/protocol.py new file mode 100644 index 000000000..766651fd5 --- /dev/null +++ b/elastica/mesh/protocol.py @@ -0,0 +1,10 @@ +from typing import Protocol + +import numpy as np +from numpy.typing import NDArray + + +class MeshProtocol(Protocol): + faces: NDArray[np.floating] + face_centers: NDArray[np.floating] + face_normals: NDArray[np.floating] diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index 5ffdcfaaf..a78425e2e 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -1,10 +1,8 @@ -from typing import Protocol, Generator, TypeVar, Any, Type, Sized +from typing import Protocol, Generator, TypeVar, Any, Type, MutableSequence from typing_extensions import Self # 3.11: from typing import Self from abc import abstractmethod -pass - from elastica.typing import ( SystemIdxType, OperatorType, @@ -35,7 +33,7 @@ def instantiate(self, *args: Any, **kwargs: Any) -> M: ... def id(self) -> Any: ... -class SystemCollectionProtocol(Sized, Protocol): +class SystemCollectionProtocol(MutableSequence, Protocol): _systems: list[SystemType] @property diff --git a/elastica/rigidbody/data_structures.py b/elastica/rigidbody/data_structures.py index d6edb4f22..95944f759 100644 --- a/elastica/rigidbody/data_structures.py +++ b/elastica/rigidbody/data_structures.py @@ -1,7 +1,8 @@ __doc__ = "Data structure wrapper for rod components" from elastica.rod.data_structures import _RodSymplecticStepperMixin -from typing import Any + +pass """ # FIXME : Explicit Stepper doesn't work as States lose the @@ -40,14 +41,5 @@ def __call__(self, time, *args, **kwargs): return self.__deriv_state """ - -class _RigidRodSymplecticStepperMixin(_RodSymplecticStepperMixin): - def __init__(self) -> None: - super(_RigidRodSymplecticStepperMixin, self).__init__() - # Expose rate returning functions in the interface - # to be used by the time-stepping algorithm - # dynamic rates needs to call update_accelerations and henc - # is another function - - def update_internal_forces_and_torques(self, *args: Any, **kwargs: Any) -> None: - pass +# TODO: Temporary solution as the structure for RigidBody is similar to Rod +_RigidRodSymplecticStepperMixin = _RodSymplecticStepperMixin diff --git a/elastica/rigidbody/mesh_rigid_body.py b/elastica/rigidbody/mesh_rigid_body.py index 4a026f9a8..f83b11c54 100644 --- a/elastica/rigidbody/mesh_rigid_body.py +++ b/elastica/rigidbody/mesh_rigid_body.py @@ -1,5 +1,8 @@ __doc__ = """rigid body class based on mesh""" +from numpy.typing import NDArray +from elastica.typing import MeshType + import numpy as np import numba from elastica._linalg import _batch_cross, _batch_norm @@ -10,12 +13,12 @@ class MeshRigidBody(RigidBodyBase): def __init__( self, - mesh, - center_of_mass, - mass_second_moment_of_inertia, - density, - volume, - ): + mesh: MeshType, + center_of_mass: NDArray[np.floating], + mass_second_moment_of_inertia: NDArray[np.floating], + density: np.floating, + volume: np.floating, + ) -> None: """ Mesh rigid body initializer. @@ -37,7 +40,7 @@ def __init__( """ # rigid body does not have elements it only have one node. We are setting n_elems to # zero for only make code to work. _bootstrap_from_data requires n_elems to be defined - self.n_elems = 1 # center_mass + self.n_elems: int = 1 # center_mass self.density = density self.volume = volume @@ -112,7 +115,7 @@ def __init__( MaxDimension.value(), 1 ) - def update_faces(self): + def update_faces(self) -> None: _update_faces( self.director_collection, self.face_centers, @@ -128,7 +131,7 @@ def update_faces(self): ) -@numba.njit(cache=True) +@numba.njit(cache=True) # type: ignore def _update_faces( director_collection, face_centers, diff --git a/elastica/typing.py b/elastica/typing.py index c23be3ee4..bf399e9d0 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -42,6 +42,8 @@ from .external_forces import NoForces from .joint import FreeJoint + from .mesh.protocol import MeshProtocol + SystemType: TypeAlias = "SystemProtocol" SystemIdxType: TypeAlias = int @@ -75,6 +77,7 @@ # Operators in elastica.modules OperatorParam = ParamSpec("OperatorParam") -OperatorType: TypeAlias = Callable[OperatorParam, None] OperatorCallbackType: TypeAlias = Callable[..., None] OperatorFinalizeType: TypeAlias = Callable[..., None] + +MeshType: TypeAlias = "MeshProtocol" diff --git a/elastica/utils.py b/elastica/utils.py index dde1e869f..429ae63ab 100644 --- a/elastica/utils.py +++ b/elastica/utils.py @@ -72,7 +72,7 @@ def value() -> Literal[3]: class Tolerance: @staticmethod - def atol() -> np.float64: + def atol() -> float: """ Static absolute tolerance method @@ -80,10 +80,10 @@ def atol() -> np.float64: ------- atol : library-wide set absolute tolerance for kernels """ - return finfo(float64).eps * 1e4 + return float(finfo(float64).eps * 1e4) @staticmethod - def rtol() -> np.float64: + def rtol() -> float: """ Static relative tolerance method @@ -91,7 +91,7 @@ def rtol() -> np.float64: ------- tol : library-wide set relative tolerance for kernels """ - return finfo(float64).eps * 1e11 + return float(finfo(float64).eps * 1e11) def perm_parity(lst: list[int]) -> int: From 555340001992002eadd589e566110ab90f226003 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 16 Jun 2024 16:11:13 -0500 Subject: [PATCH 099/134] resolve liskov error --- elastica/dissipation.py | 13 ++++---- elastica/experimental/interaction.py | 4 +-- elastica/external_forces.py | 30 +++++++------------ elastica/interaction.py | 44 ++++++---------------------- elastica/modules/damping.py | 4 +-- elastica/modules/protocol.py | 4 +-- elastica/typing.py | 6 ++-- tests/test_interaction.py | 40 ------------------------- 8 files changed, 35 insertions(+), 110 deletions(-) diff --git a/elastica/dissipation.py b/elastica/dissipation.py index 7592c85b4..2cc532918 100644 --- a/elastica/dissipation.py +++ b/elastica/dissipation.py @@ -5,7 +5,7 @@ """ from abc import ABC, abstractmethod -from typing import Any +from typing import Any, Generic, TypeVar from elastica.typing import RodType, SystemType @@ -15,7 +15,10 @@ from numpy.typing import NDArray -class DamperBase(ABC): +T = TypeVar("T") + + +class DamperBase(Generic[T], ABC): """Base class for damping module implementations. Notes @@ -29,7 +32,7 @@ class DamperBase(ABC): """ - _system: RodType + _system: T # TODO typing can be made better def __init__(self, *args: Any, **kwargs: Any) -> None: @@ -43,7 +46,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: ) @property - def system(self) -> RodType: + def system(self) -> T: """ get system (rod or rigid body) reference @@ -55,7 +58,7 @@ def system(self) -> RodType: return self._system @abstractmethod - def dampen_rates(self, system: RodType, time: np.floating) -> None: + def dampen_rates(self, system: T, time: np.floating) -> None: # TODO: In the future, we can remove rod and use self.system """ Dampen rates (velocity and/or omega) of a rod object. diff --git a/elastica/experimental/interaction.py b/elastica/experimental/interaction.py index 5609eba38..2a328934a 100644 --- a/elastica/experimental/interaction.py +++ b/elastica/experimental/interaction.py @@ -4,8 +4,8 @@ import numpy as np from elastica.external_forces import NoForces from elastica.contact_utils import _find_slipping_elements +from elastica._contact_functions import _calculate_contact_forces_cylinder_plane from elastica.interaction import ( - apply_normal_force_numba_rigid_body, InteractionPlaneRigidBody, ) @@ -93,7 +93,7 @@ def anisotropic_friction_numba_rigid_body( ( plane_response_force_mag, no_contact_point_idx, - ) = apply_normal_force_numba_rigid_body( + ) = _calculate_contact_forces_cylinder_plane( plane_origin, plane_normal, surface_tol, diff --git a/elastica/external_forces.py b/elastica/external_forces.py index 1b4124be0..dd8679db7 100644 --- a/elastica/external_forces.py +++ b/elastica/external_forces.py @@ -34,7 +34,7 @@ def __init__(self) -> None: """ pass - def apply_forces(self, system: S, time: np.floating = np.float64(0.0)) -> None: + def apply_forces(self, system: S, time: float = 0.0) -> None: """Apply forces to a rod-like object. In NoForces class, this routine simply passes. @@ -49,7 +49,7 @@ def apply_forces(self, system: S, time: np.floating = np.float64(0.0)) -> None: """ pass - def apply_torques(self, system: S, time: np.floating = np.float64(0.0)) -> None: + def apply_torques(self, system: S, time: float = 0.0) -> None: """Apply torques to a rod-like object. In NoForces class, this routine simply passes. @@ -91,7 +91,7 @@ def __init__( self.acc_gravity = acc_gravity def apply_forces( - self, system: "RodType | RigidBodyType", time: np.floating = np.float64(0.0) + self, system: "RodType | RigidBodyType", time: float = 0.0 ) -> None: self.compute_gravity_forces( self.acc_gravity, system.mass, system.external_forces @@ -162,9 +162,7 @@ def __init__( assert ramp_up_time > 0.0 self.ramp_up_time = ramp_up_time - def apply_forces( - self, system: SystemType, time: np.floating = np.float64(0.0) - ) -> None: + def apply_forces(self, system: SystemType, time: float = 0.0) -> None: self.compute_end_point_forces( system.external_forces, self.start_force, @@ -179,7 +177,7 @@ def compute_end_point_forces( external_forces: NDArray[np.floating], start_force: NDArray[np.floating], end_force: NDArray[np.floating], - time: np.floating, + time: float, ramp_up_time: np.floating, ) -> None: """ @@ -233,9 +231,7 @@ def __init__( super(UniformTorques, self).__init__() self.torque = torque * direction - def apply_torques( - self, system: SystemType, time: np.floating = np.float64(0.0) - ) -> None: + def apply_torques(self, system: SystemType, time: float = 0.0) -> None: n_elems = system.n_elems torque_on_one_element = ( _batch_product_i_k_to_ik(self.torque, np.ones((n_elems))) / n_elems @@ -273,9 +269,7 @@ def __init__( super(UniformForces, self).__init__() self.force = (force * direction).reshape(3, 1) - def apply_forces( - self, rod: SystemType, time: np.floating = np.float64(0.0) - ) -> None: + def apply_forces(self, rod: SystemType, time: float = 0.0) -> None: force_on_one_element = self.force / rod.n_elems rod.external_forces += force_on_one_element @@ -372,9 +366,7 @@ def __init__( else: self.my_spline = np.full_like(self.s, fill_value=1.0) - def apply_torques( - self, rod: SystemType, time: np.floating = np.float64(0.0) - ) -> None: + def apply_torques(self, rod: SystemType, time: float = 0.0) -> None: self.compute_muscle_torques( time, self.my_spline, @@ -391,7 +383,7 @@ def apply_torques( @staticmethod @njit(cache=True) # type: ignore def compute_muscle_torques( - time: np.floating, + time: float, my_spline: NDArray[np.floating], s: np.floating, angular_frequency: np.floating, @@ -540,9 +532,7 @@ def __init__( assert ramp_up_time >= 0.0 self.ramp_up_time = ramp_up_time - def apply_forces( - self, system: SystemType, time: np.floating = np.float64(0.0) - ) -> None: + def apply_forces(self, system: SystemType, time: float = 0.0) -> None: if time < self.ramp_up_time: # When time smaller than ramp up time apply the force in normal direction diff --git a/elastica/interaction.py b/elastica/interaction.py index c8f69eee5..784aab5c3 100644 --- a/elastica/interaction.py +++ b/elastica/interaction.py @@ -1,7 +1,7 @@ __doc__ = """ Numba implementation module containing interactions between a rod and its environment.""" -from typing import Any, NoReturn +pass import numpy as np from elastica.external_forces import NoForces from numba import njit @@ -22,7 +22,7 @@ # base class for interaction # only applies normal force no friction -class InteractionPlane: +class InteractionPlane(NoForces): """ The interaction plane class computes the plane reaction force on a rod-like object. For more details regarding the contact module refer to @@ -73,9 +73,7 @@ def __init__( self.plane_normal = plane_normal.reshape(3) self.surface_tol = 1e-4 - def apply_normal_force( - self, system: RodType - ) -> tuple[NDArray[np.floating], NDArray[np.intp]]: + def apply_forces(self, system: RodType, time: float = 0.0) -> None: """ In the case of contact with the plane, this function computes the plane reaction force on the element. @@ -114,7 +112,7 @@ def apply_normal_force( # head is at x[0] and forward means head to tail # same convention for kinetic and static # mu named as to which direction it opposes -class AnisotropicFrictionalPlane(InteractionPlane, NoForces): +class AnisotropicFrictionalPlane(InteractionPlane): """ This anisotropic friction plane class is for computing anisotropic friction forces on rods. @@ -192,7 +190,7 @@ def __init__( # kinetic and static friction should separate functions # for now putting them together to figure out common variables def apply_forces( - self, system: "RodType | RigidBodyType", time: np.floating = np.float64(0) + self, system: "RodType | RigidBodyType", time: float = 0.0 ) -> None: """ Call numba implementation to apply friction forces @@ -400,9 +398,7 @@ def __init__(self, dynamic_viscosity: np.floating) -> None: super(SlenderBodyTheory, self).__init__() self.dynamic_viscosity = dynamic_viscosity - def apply_forces( - self, system: RodType, time: np.floating = np.float64(0.0) - ) -> None: + def apply_forces(self, system: RodType, time: float = 0.0) -> None: """ This function applies hydrodynamic forces on body using the slender body theory given in @@ -427,7 +423,7 @@ def apply_forces( # base class for interaction # only applies normal force no friction -class InteractionPlaneRigidBody(InteractionPlane): +class InteractionPlaneRigidBody(NoForces): def __init__( self, k: np.floating, @@ -441,9 +437,7 @@ def __init__( self.plane_normal = plane_normal.reshape(3) self.surface_tol = 1e-4 - def apply_normal_force( - self, system: RigidBodyType - ) -> tuple[NDArray[np.floating], NDArray[np.intp]]: + def apply_forces(self, system: RigidBodyType, time: float = 0.0) -> None: """ This function computes the plane force response on the rigid body, in the case of contact. Contact model given in Eqn 4.8 Gazzola et. al. RSoS 2018 paper @@ -456,7 +450,7 @@ def apply_normal_force( ------- magnitude of the plane response """ - return _calculate_contact_forces_cylinder_plane( + _calculate_contact_forces_cylinder_plane( self.plane_origin, self.plane_normal, self.surface_tol, @@ -467,23 +461,3 @@ def apply_normal_force( system.velocity_collection, system.external_forces, ) - - -@njit(cache=True) # type: ignore -def apply_normal_force_numba_rigid_body( - plane_origin: Any, - plane_normal: Any, - surface_tol: Any, - k: Any, - nu: Any, - length: Any, - position_collection: Any, - velocity_collection: Any, - external_forces: Any, -) -> NoReturn: - - raise NotImplementedError( - "This function is removed in v0.3.2. For cylinder plane contact please use: \n" - "elastica._contact_functions._calculate_contact_forces_cylinder_plane() \n" - "For detail, refer to issue #113." - ) diff --git a/elastica/modules/damping.py b/elastica/modules/damping.py index 4fa0cf055..a3a2615e5 100644 --- a/elastica/modules/damping.py +++ b/elastica/modules/damping.py @@ -15,7 +15,7 @@ import numpy as np from elastica.dissipation import DamperBase -from elastica.typing import SystemType, SystemIdxType +from elastica.typing import RodType, SystemType, SystemIdxType from .protocol import SystemCollectionProtocol, ModuleProtocol @@ -37,7 +37,7 @@ def __init__(self: SystemCollectionProtocol) -> None: self._feature_group_constrain_rates.append(self._dampen_rates) self._feature_group_finalize.append(self._finalize_dampers) - def dampen(self: SystemCollectionProtocol, system: SystemType) -> ModuleProtocol: + def dampen(self: SystemCollectionProtocol, system: RodType) -> ModuleProtocol: """ This method applies damping on relevant user-defined system or rod-like object. You must input the system or rod-like diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index a78425e2e..9d3b51495 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -1,4 +1,4 @@ -from typing import Protocol, Generator, TypeVar, Any, Type, MutableSequence +from typing import Protocol, Generator, TypeVar, Any, Type, Sized from typing_extensions import Self # 3.11: from typing import Self from abc import abstractmethod @@ -33,7 +33,7 @@ def instantiate(self, *args: Any, **kwargs: Any) -> M: ... def id(self) -> Any: ... -class SystemCollectionProtocol(MutableSequence, Protocol): +class SystemCollectionProtocol(Sized, Protocol): _systems: list[SystemType] @property diff --git a/elastica/typing.py b/elastica/typing.py index bf399e9d0..3ebbc3fd7 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -30,9 +30,7 @@ SymplecticStepperProtocol, MemoryProtocol, ) - from memory_block.protocol import ( - BlockCosseratRodProtocol, - ) # , BlockRigidBodyProtocol + from memory_block.protocol import BlockCosseratRodProtocol, BlockRigidBodyProtocol # Modules Base Classes from .boundary_conditions import FreeBC @@ -66,7 +64,7 @@ SurfaceType: TypeAlias = "SurfaceBase" SystemCollectionType: TypeAlias = "SystemCollectionProtocol" -BlockType: TypeAlias = "BlockCosseratRodProtocol" # | "BlockRigidBodyProtocol" +BlockType: TypeAlias = "BlockCosseratRodProtocol | BlockRigidBodyProtocol" # Indexing types # TODO: Maybe just use slice?? diff --git a/tests/test_interaction.py b/tests/test_interaction.py index 4a71053f7..0267cd2a1 100644 --- a/tests/test_interaction.py +++ b/tests/test_interaction.py @@ -8,7 +8,6 @@ InteractionPlane, AnisotropicFrictionalPlane, SlenderBodyTheory, - apply_normal_force_numba_rigid_body, ) from elastica.contact_utils import ( _node_to_element_mass_or_force, @@ -853,42 +852,3 @@ def test_slender_body_matrix_product_only_xy(self, n_elem): slender_body_theory.apply_forces(rod) assert_allclose(correct_forces, rod.external_forces, atol=Tolerance.atol()) - - -@pytest.mark.parametrize("n_elem", [2, 3, 5, 10, 20]) -def test_apply_normal_force_numba_rigid_body_error_message(n_elem): - """ - This function _elements_to_nodes_inplace. We are - converting node velocities to element velocities. Here also - we are using numba to speed up the process. - - Parameters - ---------- - n_elem - - Returns - ------- - - """ - - position_collection = np.zeros((3, n_elem + 1)) - position_collection[0, :] = np.linspace(0, 1.0, n_elem + 1) - - error_message = ( - "This function is removed in v0.3.2. For cylinder plane contact please use: \n" - "elastica._contact_functions._calculate_contact_forces_cylinder_plane() \n" - "For detail, refer to issue #113." - ) - with pytest.raises(NotImplementedError) as error_info: - apply_normal_force_numba_rigid_body( - plane_origin=np.array([0.0, 0.0, 0.0]), - plane_normal=np.array([0.0, 0.0, 1.0]), - surface_tol=1e-4, - k=1.0, - nu=1.0, - length=1.0, - position_collection=position_collection, - velocity_collection=np.zeros((3, n_elem + 1)), - external_forces=np.zeros((3, n_elem + 1)), - ) - assert error_info.value.args[0] == error_message From 77952bb27698cfa01589346a67d41729b378ddad Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 16 Jun 2024 21:52:51 -0500 Subject: [PATCH 100/134] wip: finalizing typing system and system collection relation --- .../memory_block/memory_block_rigid_body.py | 16 ++++++------- elastica/memory_block/protocol.py | 24 +++++++++++++++++-- elastica/modules/base_system.py | 22 ++++++++++------- elastica/modules/constraints.py | 2 +- elastica/modules/memory_block.py | 9 ++++--- elastica/modules/protocol.py | 14 ++++++++--- elastica/systems/analytical.py | 7 +++--- elastica/timestepper/symplectic_steppers.py | 10 ++++---- elastica/typing.py | 7 +++++- pyproject.toml | 1 + tests/test_interaction.py | 14 +++++------ 11 files changed, 82 insertions(+), 44 deletions(-) diff --git a/elastica/memory_block/memory_block_rigid_body.py b/elastica/memory_block/memory_block_rigid_body.py index 0ec9d4256..788597e66 100644 --- a/elastica/memory_block/memory_block_rigid_body.py +++ b/elastica/memory_block/memory_block_rigid_body.py @@ -1,7 +1,7 @@ __doc__ = """Create block-structure class for collection of rigid body systems.""" from typing import Literal import numpy as np -from elastica.typing import SystemType, SystemIdxType +from elastica.typing import SystemIdxType, RigidBodyType from elastica.rigidbody import RigidBodyBase from elastica.rigidbody.data_structures import _RigidRodSymplecticStepperMixin @@ -9,7 +9,7 @@ class MemoryBlockRigidBody(RigidBodyBase, _RigidRodSymplecticStepperMixin): def __init__( - self, systems: list[RigidBodyBase], system_idx_list: list[SystemIdxType] + self, systems: list[RigidBodyType], system_idx_list: list[SystemIdxType] ) -> None: self.n_bodies = len(systems) @@ -26,7 +26,7 @@ def __init__( # Initialize the mixin class for symplectic time-stepper. _RigidRodSymplecticStepperMixin.__init__(self) - def _allocate_block_variables_scalars(self, systems: list[RigidBodyBase]) -> None: + def _allocate_block_variables_scalars(self, systems: list[RigidBodyType]) -> None: """ This function takes system collection and allocates the variables for block-structure and references allocated variables back to the systems. @@ -61,7 +61,7 @@ def _allocate_block_variables_scalars(self, systems: list[RigidBodyBase]) -> Non value_type="scalar", ) - def _allocate_block_variables_vectors(self, systems: list[RigidBodyBase]) -> None: + def _allocate_block_variables_vectors(self, systems: list[RigidBodyType]) -> None: """ This function takes system collection and allocates the vector variables for block-structure and references allocated vector variables back to the systems. @@ -97,7 +97,7 @@ def _allocate_block_variables_vectors(self, systems: list[RigidBodyBase]) -> Non value_type="vector", ) - def _allocate_block_variables_matrix(self, systems: list[RigidBodyBase]) -> None: + def _allocate_block_variables_matrix(self, systems: list[RigidBodyType]) -> None: """ This function takes system collection and allocates the matrix variables for block-structure and references allocated matrix variables back to the systems. @@ -134,7 +134,7 @@ def _allocate_block_variables_matrix(self, systems: list[RigidBodyBase]) -> None ) def _allocate_block_variables_for_symplectic_stepper( - self, systems: list[RigidBodyBase] + self, systems: list[RigidBodyType] ) -> None: """ This function takes system collection and allocates the variables used by symplectic @@ -181,7 +181,7 @@ def _allocate_block_variables_for_symplectic_stepper( def _map_system_properties_to_block_memory( self, mapping_dict: dict, - systems: list[RigidBodyBase], + systems: list[RigidBodyType], block_memory: np.ndarray, value_type: Literal["scalar", "vector", "tensor"], ) -> None: @@ -191,7 +191,7 @@ def _map_system_properties_to_block_memory( ---------- mapping_dict: dict Dictionary with attribute names as keys and block row index as values. - systems: list[RigidBodyBase] + systems: list[RigidBodyType] A sequence containing rigid body objects to map from. block_memory: ndarray Memory block that, at the end of the method execution, contains all designated diff --git a/elastica/memory_block/protocol.py b/elastica/memory_block/protocol.py index 80d9c7a98..c0350e87b 100644 --- a/elastica/memory_block/protocol.py +++ b/elastica/memory_block/protocol.py @@ -1,13 +1,33 @@ from typing import Protocol +from elastica.typing import SystemType + +import numpy as np from elastica.rod.protocol import CosseratRodProtocol from elastica.rigid_body.protocol import RigidBodyProtocol from elastica.systems.protocol import SymplecticSystemProtocol -class BlockCosseratRodProtocol(CosseratRodProtocol, SymplecticSystemProtocol, Protocol): +class BlockSystemProtocol(SystemType, Protocol): + @property + def n_bodies(self) -> int: + """Number of systems in the block.""" + + +class BlockOperatorsProtocol(Protocol): + + def update_internal_forces_and_torques(self, time: np.floating) -> None: ... + + def reset_external_forces_and_torques(self, time: np.floating) -> None: ... + + +class BlockCosseratRodProtocol( + CosseratRodProtocol, SymplecticSystemProtocol, BlockOperatorsProtocol, Protocol +): pass -class BlockRigidBodyProtocol(RigidBodyProtocol, SymplecticSystemProtocol, Protocol): +class BlockRigidBodyProtocol( + RigidBodyProtocol, SymplecticSystemProtocol, BlockOperatorsProtocol, Protocol +): pass diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index 74c1439bb..c31581d66 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -5,7 +5,7 @@ Basic coordinating for multiple, smaller systems that have an independently integrable interface (i.e. works with symplectic or explicit routines `timestepper.py`.) """ -from typing import Type, Generator, Iterable, Any +from typing import Type, Generator, Iterable, Any, overload from typing import final from elastica.typing import ( SystemType, @@ -69,7 +69,7 @@ def __init__(self) -> None: # List of systems to be integrated self._systems: list[SystemType] = [] - self._memory_blocks: list[BlockType] = [] + self.__final_systems: list[SystemType] = [] # Flag Finalize: Finalizing twice will cause an error, # but the error message is very misleading @@ -101,6 +101,12 @@ def _check_type(self, sys_to_be_added: Any) -> bool: def __len__(self) -> int: return len(self._systems) + @overload + def __getitem__(self, idx: int, /) -> SystemType: ... + + @overload + def __getitem__(self, idx: slice, /) -> list[SystemType]: ... + def __getitem__(self, idx, /): # type: ignore return self._systems[idx] @@ -159,9 +165,9 @@ def _get_sys_idx_if_valid(self, sys_to_be_added: SystemType) -> SystemIdxType: return sys_idx @final - def blocks(self) -> Generator[SystemType, None, None]: + def systems(self) -> Generator[SystemType, None, None]: # assert self._finalize_flag, "The simulator is not finalized." - for block in self._memory_blocks: + for block in self.__final_systems: yield block @final @@ -177,10 +183,10 @@ def finalize(self) -> None: self._finalize_flag = True # construct memory block - self._memory_blocks = construct_memory_block_structures(self._systems) - for block in self.blocks(): - # append the memory block to the simulation as a system. Memory block is the final system in the simulation. - self.append(block) + self.__final_systems = construct_memory_block_structures(self._systems) + # TODO: try to remove the _systems list for memory optimization + # self._systems.clear() + # del self._systems # Recurrent call finalize functions for all components. for finalize in self._feature_group_finalize: diff --git a/elastica/modules/constraints.py b/elastica/modules/constraints.py index 9c70e640e..0f3635383 100644 --- a/elastica/modules/constraints.py +++ b/elastica/modules/constraints.py @@ -65,7 +65,7 @@ def _finalize_constraints(self: SystemCollectionProtocol) -> None: """ from elastica._synchronize_periodic_boundary import _ConstrainPeriodicBoundaries - for block in self.blocks(): + for block in self.systems(): # append the memory block to the simulation as a system. Memory block is the final system in the simulation. if hasattr(block, "ring_rod_flag"): # Apply the constrain to synchronize the periodic boundaries of the memory rod. Find the memory block diff --git a/elastica/modules/memory_block.py b/elastica/modules/memory_block.py index f13760d88..739862c1a 100644 --- a/elastica/modules/memory_block.py +++ b/elastica/modules/memory_block.py @@ -2,8 +2,7 @@ This function is a module to construct memory blocks for different types of systems, such as Cosserat Rods, Rigid Body etc. """ - -from elastica.typing import SystemType, SystemIdxType +from elastica.typing import RodType, RigidBodyType, SystemType, SystemIdxType from elastica.rod.rod_base import RodBase from elastica.rigidbody.rigid_body import RigidBodyBase @@ -22,9 +21,9 @@ def construct_memory_block_structures(systems: list[SystemType]) -> list[SystemT ------- """ - _memory_blocks: list[SystemType] = [] - temp_list_for_cosserat_rod_systems: list[RodBase] = [] - temp_list_for_rigid_body_systems: list[RigidBodyBase] = [] + _memory_blocks: list[BlockSystemType] = [] + temp_list_for_cosserat_rod_systems: list[SystemType] = [] + temp_list_for_rigid_body_systems: list[SystemType] = [] temp_list_for_cosserat_rod_systems_idx: list[SystemIdxType] = [] temp_list_for_rigid_body_systems_idx: list[SystemIdxType] = [] diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index 9d3b51495..b65419158 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -1,4 +1,4 @@ -from typing import Protocol, Generator, TypeVar, Any, Type, Sized +from typing import Protocol, Generator, TypeVar, Any, Type, overload from typing_extensions import Self # 3.11: from typing import Self from abc import abstractmethod @@ -33,9 +33,17 @@ def instantiate(self, *args: Any, **kwargs: Any) -> M: ... def id(self) -> Any: ... -class SystemCollectionProtocol(Sized, Protocol): +class SystemCollectionProtocol(Protocol): _systems: list[SystemType] + def __len__(self) -> int: ... + + @overload + def __getitem__(self, i: slice) -> list[SystemType]: ... + @overload + def __getitem__(self, i: int) -> SystemType: ... + def __getitem__(self, i: slice | int) -> "list[SystemType] | SystemType": ... + @property def _feature_group_synchronize(self) -> OperatorGroupFIFO: ... @@ -59,7 +67,7 @@ def apply_callbacks(self, time: np.floating, current_step: int) -> None: ... @property def _feature_group_finalize(self) -> list[OperatorFinalizeType]: ... - def blocks(self) -> Generator[BlockType, None, None]: ... + def systems(self) -> Generator[SystemType, None, None]: ... def _get_sys_idx_if_valid(self, sys_to_be_added: SystemType) -> SystemIdxType: ... diff --git a/elastica/systems/analytical.py b/elastica/systems/analytical.py index fba1b4a1b..d12cea614 100644 --- a/elastica/systems/analytical.py +++ b/elastica/systems/analytical.py @@ -300,9 +300,8 @@ class CollectiveSystem: def __init__(self): self._memory_blocks = [] - self.systems = self._memory_blocks - def blocks(self): + def systems(self): return self._memory_blocks def __getitem__(self, idx): @@ -347,8 +346,8 @@ def __init__(self): super( ScalarExponentialDampedHarmonicOscillatorCollectiveSystem, self ).__init__() - self.systems.append(ScalarExponentialDecaySystem()) - self.systems.append(DampedSimpleHarmonicOscillatorSystem()) + self._memory_blocks.append(ScalarExponentialDecaySystem()) + self._memory_blocks.append(DampedSimpleHarmonicOscillatorSystem()) def make_simple_system_with_positions_directors( diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index 11184d1df..8e07eb0c2 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -95,7 +95,7 @@ def do_step( """ for kin_prefactor, kin_step, dyn_step in steps_and_prefactors[:-1]: - for system in SystemCollection.blocks(): + for system in SystemCollection.systems(): kin_step(system, time, dt) time += kin_prefactor(dt) @@ -104,14 +104,14 @@ def do_step( SystemCollection.constrain_values(time) # We need internal forces and torques because they are used by interaction module. - for system in SystemCollection.blocks(): + for system in SystemCollection.systems(): system.update_internal_forces_and_torques(time) # system.update_internal_forces_and_torques() # Add external forces, controls etc. SystemCollection.synchronize(time) - for system in SystemCollection.blocks(): + for system in SystemCollection.systems(): dyn_step(system, time, dt) # Constrain only rates @@ -121,7 +121,7 @@ def do_step( last_kin_prefactor = steps_and_prefactors[-1][0] last_kin_step = steps_and_prefactors[-1][1] - for system in SystemCollection.blocks(): + for system in SystemCollection.systems(): last_kin_step(system, time, dt) time += last_kin_prefactor(dt) SystemCollection.constrain_values(time) @@ -130,7 +130,7 @@ def do_step( SystemCollection.apply_callbacks(time, round(time / dt)) # Zero out the external forces and torques - for system in SystemCollection.blocks(): + for system in SystemCollection.systems(): system.reset_external_forces_and_torques(time) return time diff --git a/elastica/typing.py b/elastica/typing.py index 3ebbc3fd7..d13e927dc 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -30,7 +30,11 @@ SymplecticStepperProtocol, MemoryProtocol, ) - from memory_block.protocol import BlockCosseratRodProtocol, BlockRigidBodyProtocol + from .memory_block.protocol import ( + BlockCosseratRodProtocol, + BlockRigidBodyProtocol, + BlockSystemProtocol, + ) # Modules Base Classes from .boundary_conditions import FreeBC @@ -44,6 +48,7 @@ SystemType: TypeAlias = "SystemProtocol" +BlockSystemType: TypeAlias = "BlockSystemProtocol" SystemIdxType: TypeAlias = int # ModuleObjectTypes: TypeAlias = ( diff --git a/pyproject.toml b/pyproject.toml index 71b03f98d..2628c4f9a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -145,6 +145,7 @@ warn_unused_ignores = false exclude = [ "elastica/systems/analytical.py", "elastica/experimental/*", + "elastica/mesh/mesh_initializer.py", ] [tool.coverage.report] diff --git a/tests/test_interaction.py b/tests/test_interaction.py index 0267cd2a1..9d30b8067 100644 --- a/tests/test_interaction.py +++ b/tests/test_interaction.py @@ -100,7 +100,7 @@ def test_interaction_without_contact(self, n_elem): [rod, interaction_plane, external_forces] = self.initializer(n_elem, shift) - interaction_plane.apply_normal_force(rod) + interaction_plane.apply_forces(rod) correct_forces = external_forces # since no contact assert_allclose(correct_forces, rod.external_forces, atol=Tolerance.atol()) @@ -122,7 +122,7 @@ def test_interaction_plane_without_k_and_nu(self, n_elem): [rod, interaction_plane, external_forces] = self.initializer(n_elem) - interaction_plane.apply_normal_force(rod) + interaction_plane.apply_forces(rod) correct_forces = np.zeros((3, n_elem + 1)) assert_allclose(correct_forces, rod.external_forces, atol=Tolerance.atol()) @@ -153,7 +153,7 @@ def test_interaction_plane_with_k_without_nu(self, n_elem, k_w): correct_forces[..., 0] *= 0.5 correct_forces[..., -1] *= 0.5 - interaction_plane.apply_normal_force(rod) + interaction_plane.apply_forces(rod) assert_allclose(correct_forces, rod.external_forces, atol=Tolerance.atol()) @@ -189,7 +189,7 @@ def test_interaction_plane_without_k_with_nu(self, n_elem, nu_w): correct_forces[..., 0] *= 0.5 correct_forces[..., -1] *= 0.5 - interaction_plane.apply_normal_force(rod) + interaction_plane.apply_forces(rod) assert_allclose(correct_forces, rod.external_forces, atol=Tolerance.atol()) @@ -219,7 +219,7 @@ def test_interaction_when_rod_is_under_plane(self, n_elem): n_elem, shift=offset_of_plane_with_respect_to_rod, plane_normal=plane_normal ) - interaction_plane.apply_normal_force(rod) + interaction_plane.apply_forces(rod) correct_forces = np.zeros((3, n_elem + 1)) assert_allclose(correct_forces, rod.external_forces, atol=Tolerance.atol()) @@ -264,7 +264,7 @@ def test_interaction_when_rod_is_under_plane_with_k_without_nu(self, n_elem, k_w correct_forces[..., 0] *= 0.5 correct_forces[..., -1] *= 0.5 - interaction_plane.apply_normal_force(rod) + interaction_plane.apply_forces(rod) assert_allclose(correct_forces, rod.external_forces, atol=Tolerance.atol()) @@ -312,7 +312,7 @@ def test_interaction_when_rod_is_under_plane_without_k_with_nu(self, n_elem, nu_ correct_forces[..., 0] *= 0.5 correct_forces[..., -1] *= 0.5 - interaction_plane.apply_normal_force(rod) + interaction_plane.apply_forces(rod) assert_allclose(correct_forces, rod.external_forces, atol=Tolerance.atol()) From 6b39aeab7641d7133e263b7e81a8a7413b9b59e7 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Sun, 16 Jun 2024 22:06:04 -0500 Subject: [PATCH 101/134] refactor memory block related utility functions --- elastica/memory_block/memory_block_rod.py | 9 +- .../memory_block/memory_block_rod_base.py | 171 +----------------- elastica/memory_block/protocol.py | 15 +- elastica/memory_block/utils.py | 169 +++++++++++++++++ elastica/modules/memory_block.py | 21 ++- elastica/rod/data_structures.py | 10 - elastica/timestepper/symplectic_steppers.py | 4 +- ...ock_base.py => test_memory_block_utils.py} | 2 +- 8 files changed, 199 insertions(+), 202 deletions(-) create mode 100644 elastica/memory_block/utils.py rename tests/test_modules/{test_memory_block_base.py => test_memory_block_utils.py} (98%) diff --git a/elastica/memory_block/memory_block_rod.py b/elastica/memory_block/memory_block_rod.py index 4a1e8bd9f..7376ae690 100644 --- a/elastica/memory_block/memory_block_rod.py +++ b/elastica/memory_block/memory_block_rod.py @@ -2,10 +2,6 @@ import numpy as np from typing import Literal, Callable from elastica.typing import SystemIdxType, RodType -from elastica.memory_block.memory_block_rod_base import ( - make_block_memory_metadata, - make_block_memory_periodic_boundary_metadata, -) from elastica.rod.data_structures import _RodSymplecticStepperMixin from elastica.reset_functions_for_block_structure import _reset_scalar_ghost from elastica.rod.cosserat_rod import ( @@ -18,6 +14,11 @@ _synchronize_periodic_boundary_of_matrix_collection, ) +from .utils import ( + make_block_memory_metadata, + make_block_memory_periodic_boundary_metadata, +) + class MemoryBlockCosseratRod(CosseratRod, _RodSymplecticStepperMixin): """ diff --git a/elastica/memory_block/memory_block_rod_base.py b/elastica/memory_block/memory_block_rod_base.py index d8b91ff26..87fcdd6d4 100644 --- a/elastica/memory_block/memory_block_rod_base.py +++ b/elastica/memory_block/memory_block_rod_base.py @@ -1,169 +1,8 @@ -__doc__ = """Create block-structure class for collection of Cosserat rod systems.""" +__doc__ = """Deprecated module. Use memory_blocks.utils instead.""" import numpy as np import numpy.typing as npt - -def make_block_memory_metadata( - n_elems_in_rods: npt.NDArray[np.integer], -) -> tuple[ - npt.NDArray[np.integer], - npt.NDArray[np.integer], - npt.NDArray[np.integer], - npt.NDArray[np.integer], -]: - """ - This function, takes number of elements of each rod as a numpy array and computes, - ghost nodes, elements and voronoi element indexes and numbers and returns it. - - Parameters - ---------- - n_elems_in_rods: npt.NDArray - An integer array containing the number of elements in each of the n rod. - - Returns - ------- - n_elems_with_ghosts: int64 - Total number of elements with ghost elements included. There are two ghost elements - between each pair of two rods adjacent in memory block. - ghost_nodes_idx: ndarray - An integer array of length n - 1 containing the indices of ghost nodes in memory block. - ghost_elements_idx: npt.NDArray - An integer array of length 2 * (n - 1) containing the indices of ghost elements in memory block. - ghost_voronoi_idx: npt.NDArray - An integer array of length 2 * (n - 1) containing the indices of ghost Voronoi nodes in memory block. - """ - - n_nodes_in_rods = n_elems_in_rods + 1 - n_rods = n_elems_in_rods.shape[0] - - # Gap between two rods have one ghost node - # n_nodes_with_ghosts = np.sum(n_nodes_in_rods) + (n_rods - 1) - # Gap between two rods have two ghost elements : comes out to n_nodes_with_ghosts - 1 - n_elems_with_ghosts = np.sum(n_elems_in_rods) + 2 * (n_rods - 1) - # Gap between two rods have three ghost voronois : comes out to n_nodes_with_ghosts - 2 - # n_voronoi_with_ghosts = np.sum(n_voronois_in_rods) + 3 * (n_rods - 1) - - ghost_nodes_idx = np.cumsum(n_nodes_in_rods[:-1], dtype=np.int64) - # Add [0, 1, 2, ... n_rods-2] to the ghost_nodes idx to accommodate miscounting - ghost_nodes_idx += np.arange(0, n_rods - 1, dtype=np.int64) - - ghost_elems_idx = np.zeros((2 * (n_rods - 1),), dtype=np.int64) - ghost_elems_idx[::2] = ghost_nodes_idx - 1 - ghost_elems_idx[1::2] = ghost_nodes_idx.copy() - - ghost_voronoi_idx = np.zeros((3 * (n_rods - 1),), dtype=np.int64) - ghost_voronoi_idx[::3] = ghost_nodes_idx - 2 - ghost_voronoi_idx[1::3] = ghost_nodes_idx - 1 - ghost_voronoi_idx[2::3] = ghost_nodes_idx.copy() - - return n_elems_with_ghosts, ghost_nodes_idx, ghost_elems_idx, ghost_voronoi_idx - - -def make_block_memory_periodic_boundary_metadata( - n_elems_in_rods: npt.NDArray[np.integer], -) -> tuple[ - npt.NDArray[np.integer], - npt.NDArray[np.integer], - npt.NDArray[np.integer], - npt.NDArray[np.integer], -]: - """ - This function, takes the number of elements of ring rods and computes the periodic boundary node, - element and voronoi index. - - Parameters - ---------- - n_elems_in_rods : npt.NDArray - 1D (n_ring_rods,) array containing data with 'float' type. Elements of this array contains total number of - elements of one rod, including periodic boundary elements. - - Returns - ------- - n_elems - - periodic_boundary_node : npt.NDArray - 2D (2, n_periodic_boundary_nodes) array containing data with 'float' type. Vector containing periodic boundary - elements index. First dimension is the periodic boundary index, second dimension is the referenced cell index. - - periodic_boundary_elems_idx : npt.NDArray - 2D (2, n_periodic_boundary_elems) array containing data with 'float' type. Vector containing periodic boundary - nodes index. First dimension is the periodic boundary index, second dimension is the referenced cell index. - - periodic_boundary_voronoi_idx : npt.NDArray - 2D (2, n_periodic_boundary_voronoi) array containing data with 'float' type. Vector containing periodic boundary - voronoi index. First dimension is the periodic boundary index, second dimension is the referenced cell index. - - """ - - n_elem: npt.NDArray[np.integer] = n_elems_in_rods.copy() - n_rods = n_elems_in_rods.shape[0] - - periodic_boundary_node_idx = np.zeros((2, 3 * n_rods), dtype=np.int64) - # count ghost nodes, first rod does not have a ghost node at the start, so exclude first rod. - periodic_boundary_node_idx[0, 0::3][1:] = 1 - # This is for the first periodic node at the end - periodic_boundary_node_idx[0, 1::3] = 1 + n_elem - # This is for the second periodic node at the end - periodic_boundary_node_idx[0, 2::3] = 1 - periodic_boundary_node_idx[0, :] = np.cumsum(periodic_boundary_node_idx[0, :]) - # Add [0, 1, 2, ..., n_rods] to the periodic boundary nodes to accommodate miscounting - periodic_boundary_node_idx[0, :] += np.repeat( - np.arange(0, n_rods, dtype=np.int64), 3 - ) - # Now fill the reference node idx, to copy and correct periodic boundary nodes - # First fill with the reference node idx of the first periodic node. This is the last node of the actual rod - # (without ghost and periodic nodes). - periodic_boundary_node_idx[1, 0::3] = periodic_boundary_node_idx[0, 1::3] - 1 - # Second fill with the reference node idx of the second periodic node. This is the first node of the actual rod - # (without ghost and periodic nodes). - periodic_boundary_node_idx[1, 1::3] = periodic_boundary_node_idx[0, 0::3] + 1 - # Third fill with the reference node idx of the third periodic node. This is the second node of the actual rod - # (without ghost and periodic nodes). - periodic_boundary_node_idx[1, 2::3] = periodic_boundary_node_idx[0, 0::3] + 2 - - periodic_boundary_elems_idx = np.zeros((2, 2 * n_rods), dtype=np.int64) - # count ghost elems, first rod does not have a ghost elem at the start, so exclude first rod. - periodic_boundary_elems_idx[0, 0::2][1:] = 2 - # This is for the first periodic elem at the end - periodic_boundary_elems_idx[0, 1::2] = 1 + n_elem - periodic_boundary_elems_idx[0, :] = np.cumsum(periodic_boundary_elems_idx[0, :]) - # Add [0, 1, 2, ..., n_rods] to the periodic boundary elems to accommodate miscounting - periodic_boundary_elems_idx[0, :] += np.repeat( - np.arange(0, n_rods, dtype=np.int64), 2 - ) - # Now fill the reference element idx, to copy and correct periodic boundary elements - # First fill with the reference element idx of the first periodic element. This is the last element of the actual - # rod - # (without ghost and periodic elements). - periodic_boundary_elems_idx[1, 0::2] = periodic_boundary_elems_idx[0, 1::2] - 1 - # Second fill with the reference element idx of the second periodic element. This is the first element of the actual - # rod - # (without ghost and periodic elements). - periodic_boundary_elems_idx[1, 1::2] = periodic_boundary_elems_idx[0, 0::2] + 1 - - periodic_boundary_voronoi_idx = np.zeros((2, n_rods), dtype=np.int64) - # count ghost voronoi, first rod does not have a ghost voronoi at the start, so exclude first rod. - periodic_boundary_voronoi_idx[0, 0::1][1:] = 3 - # This is for the first periodic voronoi at the end - periodic_boundary_voronoi_idx[0, 1:] += n_elem[:-1] - periodic_boundary_voronoi_idx[0, :] = np.cumsum(periodic_boundary_voronoi_idx[0, :]) - # Add [0, 1, 2, ..., n_rods] to the periodic boundary voronoi to accommodate miscounting - periodic_boundary_voronoi_idx[0, :] += np.repeat( - np.arange(0, n_rods, dtype=np.int64), 1 - ) - # Now fill the reference voronoi idx, to copy and correct periodic boundary voronoi - # Fill with the reference voronoi idx of the periodic voronoi. This is the last voronoi of the actual rod - # (without ghost and periodic voronoi). - periodic_boundary_voronoi_idx[1, :] = ( - periodic_boundary_voronoi_idx[0, :] + n_elem[:] - ) - - # Increase the n_elem in rods by 2 because we are adding two periodic boundary elements - n_elem = n_elem + 2 - - return ( - n_elem, - periodic_boundary_node_idx, - periodic_boundary_elems_idx, - periodic_boundary_voronoi_idx, - ) +from .utils import ( + make_block_memory_metadata, + make_block_memory_periodic_boundary_metadata, +) diff --git a/elastica/memory_block/protocol.py b/elastica/memory_block/protocol.py index c0350e87b..65980ee6b 100644 --- a/elastica/memory_block/protocol.py +++ b/elastica/memory_block/protocol.py @@ -14,20 +14,9 @@ def n_bodies(self) -> int: """Number of systems in the block.""" -class BlockOperatorsProtocol(Protocol): - - def update_internal_forces_and_torques(self, time: np.floating) -> None: ... - - def reset_external_forces_and_torques(self, time: np.floating) -> None: ... - - -class BlockCosseratRodProtocol( - CosseratRodProtocol, SymplecticSystemProtocol, BlockOperatorsProtocol, Protocol -): +class BlockCosseratRodProtocol(CosseratRodProtocol, SymplecticSystemProtocol, Protocol): pass -class BlockRigidBodyProtocol( - RigidBodyProtocol, SymplecticSystemProtocol, BlockOperatorsProtocol, Protocol -): +class BlockRigidBodyProtocol(RigidBodyProtocol, SymplecticSystemProtocol, Protocol): pass diff --git a/elastica/memory_block/utils.py b/elastica/memory_block/utils.py new file mode 100644 index 000000000..d8b91ff26 --- /dev/null +++ b/elastica/memory_block/utils.py @@ -0,0 +1,169 @@ +__doc__ = """Create block-structure class for collection of Cosserat rod systems.""" +import numpy as np +import numpy.typing as npt + + +def make_block_memory_metadata( + n_elems_in_rods: npt.NDArray[np.integer], +) -> tuple[ + npt.NDArray[np.integer], + npt.NDArray[np.integer], + npt.NDArray[np.integer], + npt.NDArray[np.integer], +]: + """ + This function, takes number of elements of each rod as a numpy array and computes, + ghost nodes, elements and voronoi element indexes and numbers and returns it. + + Parameters + ---------- + n_elems_in_rods: npt.NDArray + An integer array containing the number of elements in each of the n rod. + + Returns + ------- + n_elems_with_ghosts: int64 + Total number of elements with ghost elements included. There are two ghost elements + between each pair of two rods adjacent in memory block. + ghost_nodes_idx: ndarray + An integer array of length n - 1 containing the indices of ghost nodes in memory block. + ghost_elements_idx: npt.NDArray + An integer array of length 2 * (n - 1) containing the indices of ghost elements in memory block. + ghost_voronoi_idx: npt.NDArray + An integer array of length 2 * (n - 1) containing the indices of ghost Voronoi nodes in memory block. + """ + + n_nodes_in_rods = n_elems_in_rods + 1 + n_rods = n_elems_in_rods.shape[0] + + # Gap between two rods have one ghost node + # n_nodes_with_ghosts = np.sum(n_nodes_in_rods) + (n_rods - 1) + # Gap between two rods have two ghost elements : comes out to n_nodes_with_ghosts - 1 + n_elems_with_ghosts = np.sum(n_elems_in_rods) + 2 * (n_rods - 1) + # Gap between two rods have three ghost voronois : comes out to n_nodes_with_ghosts - 2 + # n_voronoi_with_ghosts = np.sum(n_voronois_in_rods) + 3 * (n_rods - 1) + + ghost_nodes_idx = np.cumsum(n_nodes_in_rods[:-1], dtype=np.int64) + # Add [0, 1, 2, ... n_rods-2] to the ghost_nodes idx to accommodate miscounting + ghost_nodes_idx += np.arange(0, n_rods - 1, dtype=np.int64) + + ghost_elems_idx = np.zeros((2 * (n_rods - 1),), dtype=np.int64) + ghost_elems_idx[::2] = ghost_nodes_idx - 1 + ghost_elems_idx[1::2] = ghost_nodes_idx.copy() + + ghost_voronoi_idx = np.zeros((3 * (n_rods - 1),), dtype=np.int64) + ghost_voronoi_idx[::3] = ghost_nodes_idx - 2 + ghost_voronoi_idx[1::3] = ghost_nodes_idx - 1 + ghost_voronoi_idx[2::3] = ghost_nodes_idx.copy() + + return n_elems_with_ghosts, ghost_nodes_idx, ghost_elems_idx, ghost_voronoi_idx + + +def make_block_memory_periodic_boundary_metadata( + n_elems_in_rods: npt.NDArray[np.integer], +) -> tuple[ + npt.NDArray[np.integer], + npt.NDArray[np.integer], + npt.NDArray[np.integer], + npt.NDArray[np.integer], +]: + """ + This function, takes the number of elements of ring rods and computes the periodic boundary node, + element and voronoi index. + + Parameters + ---------- + n_elems_in_rods : npt.NDArray + 1D (n_ring_rods,) array containing data with 'float' type. Elements of this array contains total number of + elements of one rod, including periodic boundary elements. + + Returns + ------- + n_elems + + periodic_boundary_node : npt.NDArray + 2D (2, n_periodic_boundary_nodes) array containing data with 'float' type. Vector containing periodic boundary + elements index. First dimension is the periodic boundary index, second dimension is the referenced cell index. + + periodic_boundary_elems_idx : npt.NDArray + 2D (2, n_periodic_boundary_elems) array containing data with 'float' type. Vector containing periodic boundary + nodes index. First dimension is the periodic boundary index, second dimension is the referenced cell index. + + periodic_boundary_voronoi_idx : npt.NDArray + 2D (2, n_periodic_boundary_voronoi) array containing data with 'float' type. Vector containing periodic boundary + voronoi index. First dimension is the periodic boundary index, second dimension is the referenced cell index. + + """ + + n_elem: npt.NDArray[np.integer] = n_elems_in_rods.copy() + n_rods = n_elems_in_rods.shape[0] + + periodic_boundary_node_idx = np.zeros((2, 3 * n_rods), dtype=np.int64) + # count ghost nodes, first rod does not have a ghost node at the start, so exclude first rod. + periodic_boundary_node_idx[0, 0::3][1:] = 1 + # This is for the first periodic node at the end + periodic_boundary_node_idx[0, 1::3] = 1 + n_elem + # This is for the second periodic node at the end + periodic_boundary_node_idx[0, 2::3] = 1 + periodic_boundary_node_idx[0, :] = np.cumsum(periodic_boundary_node_idx[0, :]) + # Add [0, 1, 2, ..., n_rods] to the periodic boundary nodes to accommodate miscounting + periodic_boundary_node_idx[0, :] += np.repeat( + np.arange(0, n_rods, dtype=np.int64), 3 + ) + # Now fill the reference node idx, to copy and correct periodic boundary nodes + # First fill with the reference node idx of the first periodic node. This is the last node of the actual rod + # (without ghost and periodic nodes). + periodic_boundary_node_idx[1, 0::3] = periodic_boundary_node_idx[0, 1::3] - 1 + # Second fill with the reference node idx of the second periodic node. This is the first node of the actual rod + # (without ghost and periodic nodes). + periodic_boundary_node_idx[1, 1::3] = periodic_boundary_node_idx[0, 0::3] + 1 + # Third fill with the reference node idx of the third periodic node. This is the second node of the actual rod + # (without ghost and periodic nodes). + periodic_boundary_node_idx[1, 2::3] = periodic_boundary_node_idx[0, 0::3] + 2 + + periodic_boundary_elems_idx = np.zeros((2, 2 * n_rods), dtype=np.int64) + # count ghost elems, first rod does not have a ghost elem at the start, so exclude first rod. + periodic_boundary_elems_idx[0, 0::2][1:] = 2 + # This is for the first periodic elem at the end + periodic_boundary_elems_idx[0, 1::2] = 1 + n_elem + periodic_boundary_elems_idx[0, :] = np.cumsum(periodic_boundary_elems_idx[0, :]) + # Add [0, 1, 2, ..., n_rods] to the periodic boundary elems to accommodate miscounting + periodic_boundary_elems_idx[0, :] += np.repeat( + np.arange(0, n_rods, dtype=np.int64), 2 + ) + # Now fill the reference element idx, to copy and correct periodic boundary elements + # First fill with the reference element idx of the first periodic element. This is the last element of the actual + # rod + # (without ghost and periodic elements). + periodic_boundary_elems_idx[1, 0::2] = periodic_boundary_elems_idx[0, 1::2] - 1 + # Second fill with the reference element idx of the second periodic element. This is the first element of the actual + # rod + # (without ghost and periodic elements). + periodic_boundary_elems_idx[1, 1::2] = periodic_boundary_elems_idx[0, 0::2] + 1 + + periodic_boundary_voronoi_idx = np.zeros((2, n_rods), dtype=np.int64) + # count ghost voronoi, first rod does not have a ghost voronoi at the start, so exclude first rod. + periodic_boundary_voronoi_idx[0, 0::1][1:] = 3 + # This is for the first periodic voronoi at the end + periodic_boundary_voronoi_idx[0, 1:] += n_elem[:-1] + periodic_boundary_voronoi_idx[0, :] = np.cumsum(periodic_boundary_voronoi_idx[0, :]) + # Add [0, 1, 2, ..., n_rods] to the periodic boundary voronoi to accommodate miscounting + periodic_boundary_voronoi_idx[0, :] += np.repeat( + np.arange(0, n_rods, dtype=np.int64), 1 + ) + # Now fill the reference voronoi idx, to copy and correct periodic boundary voronoi + # Fill with the reference voronoi idx of the periodic voronoi. This is the last voronoi of the actual rod + # (without ghost and periodic voronoi). + periodic_boundary_voronoi_idx[1, :] = ( + periodic_boundary_voronoi_idx[0, :] + n_elem[:] + ) + + # Increase the n_elem in rods by 2 because we are adding two periodic boundary elements + n_elem = n_elem + 2 + + return ( + n_elem, + periodic_boundary_node_idx, + periodic_boundary_elems_idx, + periodic_boundary_voronoi_idx, + ) diff --git a/elastica/modules/memory_block.py b/elastica/modules/memory_block.py index 739862c1a..3f69b57e0 100644 --- a/elastica/modules/memory_block.py +++ b/elastica/modules/memory_block.py @@ -2,7 +2,14 @@ This function is a module to construct memory blocks for different types of systems, such as Cosserat Rods, Rigid Body etc. """ -from elastica.typing import RodType, RigidBodyType, SystemType, SystemIdxType +from typing import cast +from elastica.typing import ( + RodType, + RigidBodyType, + SystemType, + SystemIdxType, + BlockSystemType, +) from elastica.rod.rod_base import RodBase from elastica.rigidbody.rigid_body import RigidBodyBase @@ -22,19 +29,21 @@ def construct_memory_block_structures(systems: list[SystemType]) -> list[SystemT """ _memory_blocks: list[BlockSystemType] = [] - temp_list_for_cosserat_rod_systems: list[SystemType] = [] - temp_list_for_rigid_body_systems: list[SystemType] = [] + temp_list_for_cosserat_rod_systems: list[RodType] = [] + temp_list_for_rigid_body_systems: list[RigidBodyType] = [] temp_list_for_cosserat_rod_systems_idx: list[SystemIdxType] = [] temp_list_for_rigid_body_systems_idx: list[SystemIdxType] = [] for system_idx, sys_to_be_added in enumerate(systems): if isinstance(sys_to_be_added, RodBase): - temp_list_for_cosserat_rod_systems.append(sys_to_be_added) + rod_system = cast(RodType, sys_to_be_added) + temp_list_for_cosserat_rod_systems.append(rod_system) temp_list_for_cosserat_rod_systems_idx.append(system_idx) elif isinstance(sys_to_be_added, RigidBodyBase): - temp_list_for_rigid_body_systems.append(sys_to_be_added) + rigid_body_system = cast(RigidBodyType, sys_to_be_added) + temp_list_for_rigid_body_systems.append(rigid_body_system) temp_list_for_rigid_body_systems_idx.append(system_idx) elif isinstance(sys_to_be_added, SurfaceBase): @@ -71,4 +80,4 @@ def construct_memory_block_structures(systems: list[SystemType]) -> list[SystemT ) ) - return _memory_blocks + return list(_memory_blocks) diff --git a/elastica/rod/data_structures.py b/elastica/rod/data_structures.py index dcbbee20c..5db3d03df 100644 --- a/elastica/rod/data_structures.py +++ b/elastica/rod/data_structures.py @@ -67,11 +67,6 @@ def __init__(self: BlockCosseratRodProtocol) -> None: # is another function self.kinematic_rates = self.dynamic_states.kinematic_rates - def update_internal_forces_and_torques( - self: BlockCosseratRodProtocol, time: np.floating - ) -> None: - self.compute_internal_forces_and_torques(time) - def dynamic_rates( self: BlockCosseratRodProtocol, time: np.floating, @@ -80,11 +75,6 @@ def dynamic_rates( self.update_accelerations(time) return self.dynamic_states.dynamic_rates(time, prefac) - def reset_external_forces_and_torques( - self: BlockCosseratRodProtocol, time: np.floating - ) -> None: - self.zeroed_out_external_forces_and_torques(time) - def _bootstrap_from_data( stepper_type: str, diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index 8e07eb0c2..5d7c0ca37 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -105,7 +105,7 @@ def do_step( # We need internal forces and torques because they are used by interaction module. for system in SystemCollection.systems(): - system.update_internal_forces_and_torques(time) + system.compute_internal_forces_and_torques(time) # system.update_internal_forces_and_torques() # Add external forces, controls etc. @@ -131,7 +131,7 @@ def do_step( # Zero out the external forces and torques for system in SystemCollection.systems(): - system.reset_external_forces_and_torques(time) + system.zeroed_out_external_forces_and_torques(time) return time diff --git a/tests/test_modules/test_memory_block_base.py b/tests/test_modules/test_memory_block_utils.py similarity index 98% rename from tests/test_modules/test_memory_block_base.py rename to tests/test_modules/test_memory_block_utils.py index 9dde39234..580addabd 100644 --- a/tests/test_modules/test_memory_block_base.py +++ b/tests/test_modules/test_memory_block_utils.py @@ -3,7 +3,7 @@ import pytest import numpy as np from numpy.testing import assert_array_equal -from elastica.memory_block.memory_block_rod_base import ( +from elastica.memory_block.utils import ( make_block_memory_metadata, make_block_memory_periodic_boundary_metadata, ) From cdce69e0ec5f5e6ca154952b3b9198dd880d3f47 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Mon, 17 Jun 2024 01:59:00 -0500 Subject: [PATCH 102/134] type: finish organizing prototype --- elastica/boundary_conditions.py | 42 ++++++++++++++----- elastica/callback_functions.py | 6 +-- elastica/external_forces.py | 16 ++++--- .../memory_block/memory_block_rigid_body.py | 4 +- elastica/memory_block/memory_block_rod.py | 11 ++--- elastica/memory_block/protocol.py | 8 ++-- elastica/memory_block/utils.py | 10 ++--- elastica/modules/connections.py | 3 +- elastica/modules/constraints.py | 8 ++-- elastica/rigidbody/cylinder.py | 4 +- elastica/rigidbody/mesh_rigid_body.py | 2 +- elastica/rigidbody/protocol.py | 8 ++-- elastica/rigidbody/rigid_body.py | 15 ++++++- elastica/rigidbody/sphere.py | 6 +-- elastica/rod/data_structures.py | 8 ++-- elastica/rod/protocol.py | 7 +++- elastica/rod/rod_base.py | 4 ++ elastica/systems/analytical.py | 4 +- elastica/systems/protocol.py | 28 ++++++------- .../test_memory_block_rigid_body.py | 2 +- 20 files changed, 122 insertions(+), 74 deletions(-) diff --git a/elastica/boundary_conditions.py b/elastica/boundary_conditions.py index d351ed323..6b25b8653 100644 --- a/elastica/boundary_conditions.py +++ b/elastica/boundary_conditions.py @@ -12,7 +12,7 @@ from elastica._linalg import _batch_matvec, _batch_matrix_transpose from elastica._rotations import _get_rotation_matrix -from elastica.typing import SystemType, RodType +from elastica.typing import SystemType, RodType, RigidBodyType S = TypeVar("S") @@ -110,11 +110,15 @@ class FreeBC(ConstraintBase): def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) - def constrain_values(self, system: SystemType, time: np.floating) -> None: + def constrain_values( + self, system: "RodType | RigidBodyType", time: np.floating + ) -> None: """In FreeBC, this routine simply passes.""" pass - def constrain_rates(self, system: SystemType, time: np.floating) -> None: + def constrain_rates( + self, system: "RodType | RigidBodyType", time: np.floating + ) -> None: """In FreeBC, this routine simply passes.""" pass @@ -168,7 +172,9 @@ def __init__( self.fixed_position_collection = np.array(fixed_position) self.fixed_directors_collection = np.array(fixed_directors) - def constrain_values(self, system: SystemType, time: np.floating) -> None: + def constrain_values( + self, system: "RodType | RigidBodyType", time: np.floating + ) -> None: # system.position_collection[..., 0] = self.fixed_position # system.director_collection[..., 0] = self.fixed_directors self.compute_constrain_values( @@ -178,7 +184,9 @@ def constrain_values(self, system: SystemType, time: np.floating) -> None: self.fixed_directors_collection, ) - def constrain_rates(self, system: SystemType, time: np.floating) -> None: + def constrain_rates( + self, system: "RodType | RigidBodyType", time: np.floating + ) -> None: # system.velocity_collection[..., 0] = 0.0 # system.omega_collection[..., 0] = 0.0 self.compute_constrain_rates( @@ -343,7 +351,9 @@ def __init__( ) self.rotational_constraint_selector = rotational_constraint_selector.astype(int) - def constrain_values(self, system: SystemType, time: np.floating) -> None: + def constrain_values( + self, system: "RodType | RigidBodyType", time: np.floating + ) -> None: if self.constrained_position_idx.size: self.nb_constrain_translational_values( system.position_collection, @@ -352,7 +362,9 @@ def constrain_values(self, system: SystemType, time: np.floating) -> None: self.translational_constraint_selector, ) - def constrain_rates(self, system: SystemType, time: np.floating) -> None: + def constrain_rates( + self, system: "RodType | RigidBodyType", time: np.floating + ) -> None: if self.constrained_position_idx.size: self.nb_constrain_translational_rates( system.velocity_collection, @@ -528,7 +540,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: **kwargs, ) - def constrain_values(self, system: SystemType, time: np.floating) -> None: + def constrain_values( + self, system: "RodType | RigidBodyType", time: np.floating + ) -> None: if self.constrained_position_idx.size: self.nb_constrain_translational_values( system.position_collection, @@ -542,7 +556,9 @@ def constrain_values(self, system: SystemType, time: np.floating) -> None: self.constrained_director_idx, ) - def constrain_rates(self, system: SystemType, time: np.floating) -> None: + def constrain_rates( + self, system: "RodType | RigidBodyType", time: np.floating + ) -> None: if self.constrained_position_idx.size: self.nb_constrain_translational_rates( system.velocity_collection, @@ -746,7 +762,9 @@ def __init__( @ director_end ) # rotation_matrix wants vectors 3,1 - def constrain_values(self, rod: SystemType, time: np.floating) -> None: + def constrain_values( + self, rod: "RodType | RigidBodyType", time: np.floating + ) -> None: if time > self.twisting_time: rod.position_collection[..., 0] = self.final_start_position rod.position_collection[..., -1] = self.final_end_position @@ -754,7 +772,9 @@ def constrain_values(self, rod: SystemType, time: np.floating) -> None: rod.director_collection[..., 0] = self.final_start_directors rod.director_collection[..., -1] = self.final_end_directors - def constrain_rates(self, rod: SystemType, time: np.floating) -> None: + def constrain_rates( + self, rod: "RodType | RigidBodyType", time: np.floating + ) -> None: if time > self.twisting_time: rod.velocity_collection[..., 0] = 0.0 rod.omega_collection[..., 0] = 0.0 diff --git a/elastica/callback_functions.py b/elastica/callback_functions.py index cc03be68b..8cfdab9ce 100644 --- a/elastica/callback_functions.py +++ b/elastica/callback_functions.py @@ -11,7 +11,7 @@ from collections import defaultdict -from elastica.typing import RodType, SystemType +from elastica.typing import RodType, RigidBodyType, SystemType T = TypeVar("T") @@ -83,7 +83,7 @@ def __init__(self, step_skip: int, callback_params: dict) -> None: self.callback_params = callback_params def make_callback( - self, system: SystemType, time: np.floating, current_step: int + self, system: "RodType | RigidBodyType", time: np.floating, current_step: int ) -> None: if current_step % self.sample_every == 0: @@ -203,7 +203,7 @@ def __init__( self._ext = "pkl" def make_callback( - self, system: SystemType, time: np.floating, current_step: int + self, system: "RodType | RigidBodyType", time: np.floating, current_step: int ) -> None: """ diff --git a/elastica/external_forces.py b/elastica/external_forces.py index dd8679db7..ae85a3569 100644 --- a/elastica/external_forces.py +++ b/elastica/external_forces.py @@ -162,7 +162,9 @@ def __init__( assert ramp_up_time > 0.0 self.ramp_up_time = ramp_up_time - def apply_forces(self, system: SystemType, time: float = 0.0) -> None: + def apply_forces( + self, system: "RodType | RigidBodyType", time: float = 0.0 + ) -> None: self.compute_end_point_forces( system.external_forces, self.start_force, @@ -231,7 +233,9 @@ def __init__( super(UniformTorques, self).__init__() self.torque = torque * direction - def apply_torques(self, system: SystemType, time: float = 0.0) -> None: + def apply_torques( + self, system: "RodType | RigidBodyType", time: float = 0.0 + ) -> None: n_elems = system.n_elems torque_on_one_element = ( _batch_product_i_k_to_ik(self.torque, np.ones((n_elems))) / n_elems @@ -269,7 +273,7 @@ def __init__( super(UniformForces, self).__init__() self.force = (force * direction).reshape(3, 1) - def apply_forces(self, rod: SystemType, time: float = 0.0) -> None: + def apply_forces(self, rod: "RodType | RigidBodyType", time: float = 0.0) -> None: force_on_one_element = self.force / rod.n_elems rod.external_forces += force_on_one_element @@ -366,7 +370,7 @@ def __init__( else: self.my_spline = np.full_like(self.s, fill_value=1.0) - def apply_torques(self, rod: SystemType, time: float = 0.0) -> None: + def apply_torques(self, rod: "RodType | RigidBodyType", time: float = 0.0) -> None: self.compute_muscle_torques( time, self.my_spline, @@ -532,7 +536,9 @@ def __init__( assert ramp_up_time >= 0.0 self.ramp_up_time = ramp_up_time - def apply_forces(self, system: SystemType, time: float = 0.0) -> None: + def apply_forces( + self, system: "RodType | RigidBodyType", time: float = 0.0 + ) -> None: if time < self.ramp_up_time: # When time smaller than ramp up time apply the force in normal direction diff --git a/elastica/memory_block/memory_block_rigid_body.py b/elastica/memory_block/memory_block_rigid_body.py index 788597e66..38ebdb8f1 100644 --- a/elastica/memory_block/memory_block_rigid_body.py +++ b/elastica/memory_block/memory_block_rigid_body.py @@ -12,8 +12,8 @@ def __init__( self, systems: list[RigidBodyType], system_idx_list: list[SystemIdxType] ) -> None: - self.n_bodies = len(systems) - self.n_elems = self.n_bodies + self.n_systems = len(systems) + self.n_elems = self.n_systems self.n_nodes = self.n_elems self.system_idx_list = np.array(system_idx_list, dtype=np.int64) diff --git a/elastica/memory_block/memory_block_rod.py b/elastica/memory_block/memory_block_rod.py index 7376ae690..6cc07a086 100644 --- a/elastica/memory_block/memory_block_rod.py +++ b/elastica/memory_block/memory_block_rod.py @@ -33,6 +33,7 @@ class MemoryBlockCosseratRod(CosseratRod, _RodSymplecticStepperMixin): def __init__( self, systems: list[RodType], system_idx_list: list[SystemIdxType] ) -> None: + self.n_systems = len(systems) # separate straight and ring rods system_straight_rod = [] @@ -62,8 +63,8 @@ def __init__( [x.n_elems for x in system_ring_rod], dtype=np.int64 ) - n_straight_rods = len(system_straight_rod) - n_ring_rods = len(system_ring_rod) + n_straight_rods: int = len(system_straight_rod) + n_ring_rods: int = len(system_ring_rod) # self.n_elems_in_rods = np.array([x.n_elems for x in systems], dtype=np.int) self.n_elems_in_rods = np.hstack((n_elems_straight_rods, n_elems_ring_rods + 2)) @@ -133,13 +134,13 @@ def __init__( # we place first straight rods, then ring rods. # TODO: in future consider a better implementation for packing problem. # +1 is because we want to start from next idx, where periodic boundary starts - self.periodic_boundary_nodes_idx += ( + self.periodic_boundary_nodes_idx += ( # type: ignore self.ghost_nodes_idx[n_straight_rods - 1] + 1 ) - self.periodic_boundary_elems_idx += ( + self.periodic_boundary_elems_idx += ( # type: ignore self.ghost_elems_idx[1::2][n_straight_rods - 1] + 1 ) - self.periodic_boundary_voronoi_idx += ( + self.periodic_boundary_voronoi_idx += ( # type: ignore self.ghost_voronoi_idx[2::3][n_straight_rods - 1] + 1 ) diff --git a/elastica/memory_block/protocol.py b/elastica/memory_block/protocol.py index 65980ee6b..279f01580 100644 --- a/elastica/memory_block/protocol.py +++ b/elastica/memory_block/protocol.py @@ -1,16 +1,16 @@ from typing import Protocol -from elastica.typing import SystemType +from elastica.systems.protocol import SystemProtocol import numpy as np from elastica.rod.protocol import CosseratRodProtocol -from elastica.rigid_body.protocol import RigidBodyProtocol +from elastica.rigidbody.protocol import RigidBodyProtocol from elastica.systems.protocol import SymplecticSystemProtocol -class BlockSystemProtocol(SystemType, Protocol): +class BlockSystemProtocol(SystemProtocol, Protocol): @property - def n_bodies(self) -> int: + def n_systems(self) -> int: """Number of systems in the block.""" diff --git a/elastica/memory_block/utils.py b/elastica/memory_block/utils.py index d8b91ff26..c7219e405 100644 --- a/elastica/memory_block/utils.py +++ b/elastica/memory_block/utils.py @@ -6,7 +6,7 @@ def make_block_memory_metadata( n_elems_in_rods: npt.NDArray[np.integer], ) -> tuple[ - npt.NDArray[np.integer], + int, npt.NDArray[np.integer], npt.NDArray[np.integer], npt.NDArray[np.integer], @@ -74,7 +74,7 @@ def make_block_memory_periodic_boundary_metadata( Parameters ---------- n_elems_in_rods : npt.NDArray - 1D (n_ring_rods,) array containing data with 'float' type. Elements of this array contains total number of + 1D (n_ring_rods,) array containing data with 'int' type. Elements of this array contains total number of elements of one rod, including periodic boundary elements. Returns @@ -82,15 +82,15 @@ def make_block_memory_periodic_boundary_metadata( n_elems periodic_boundary_node : npt.NDArray - 2D (2, n_periodic_boundary_nodes) array containing data with 'float' type. Vector containing periodic boundary + 2D (2, n_periodic_boundary_nodes) array containing data with 'int' type. Vector containing periodic boundary elements index. First dimension is the periodic boundary index, second dimension is the referenced cell index. periodic_boundary_elems_idx : npt.NDArray - 2D (2, n_periodic_boundary_elems) array containing data with 'float' type. Vector containing periodic boundary + 2D (2, n_periodic_boundary_elems) array containing data with 'int' type. Vector containing periodic boundary nodes index. First dimension is the periodic boundary index, second dimension is the referenced cell index. periodic_boundary_voronoi_idx : npt.NDArray - 2D (2, n_periodic_boundary_voronoi) array containing data with 'float' type. Vector containing periodic boundary + 2D (2, n_periodic_boundary_voronoi) array containing data with 'int' type. Vector containing periodic boundary voronoi index. First dimension is the periodic boundary index, second dimension is the referenced cell index. """ diff --git a/elastica/modules/connections.py b/elastica/modules/connections.py index 2b82bb7ce..d90173853 100644 --- a/elastica/modules/connections.py +++ b/elastica/modules/connections.py @@ -65,10 +65,9 @@ def connect( ------- """ + # For each system identified, get max dofs sys_idx_first = self._get_sys_idx_if_valid(first_rod) sys_idx_second = self._get_sys_idx_if_valid(second_rod) - - # For each system identified, get max dofs sys_dofs_first = self._systems[sys_idx_first].n_elems sys_dofs_second = self._systems[sys_idx_second].n_elems diff --git a/elastica/modules/constraints.py b/elastica/modules/constraints.py index 0f3635383..bfa452124 100644 --- a/elastica/modules/constraints.py +++ b/elastica/modules/constraints.py @@ -11,7 +11,7 @@ from elastica.boundary_conditions import ConstraintBase -from elastica.typing import SystemType, SystemIdxType, ConstrainingIndex +from elastica.typing import SystemIdxType, ConstrainingIndex, RigidBodyType, RodType from .protocol import SystemCollectionProtocol, ModuleProtocol @@ -34,7 +34,9 @@ def __init__(self: SystemCollectionProtocol) -> None: self._feature_group_constrain_rates.append(self._constrain_rates) self._feature_group_finalize.append(self._finalize_constraints) - def constrain(self: SystemCollectionProtocol, system: SystemType) -> ModuleProtocol: + def constrain( + self: SystemCollectionProtocol, system: "RodType | RigidBodyType" + ) -> ModuleProtocol: """ This method enforces a displacement boundary conditions to the relevant user-defined system or rod-like object. You must input the system or rod-like @@ -178,7 +180,7 @@ def using( def id(self) -> SystemIdxType: return self._sys_idx - def instantiate(self, system: SystemType) -> ConstraintBase: + def instantiate(self, system: "RodType | RigidBodyType") -> ConstraintBase: """Constructs a constraint after checks""" if not hasattr(self, "_bc_cls"): raise RuntimeError( diff --git a/elastica/rigidbody/cylinder.py b/elastica/rigidbody/cylinder.py index 079ce7d53..053dfdee3 100644 --- a/elastica/rigidbody/cylinder.py +++ b/elastica/rigidbody/cylinder.py @@ -72,8 +72,8 @@ def _check_lower_bound( dim: int_t = MaxDimension.value() # This is for a rigid body cylinder - self.volume = np.pi * base_radius * base_radius * base_length - self.mass = np.array([self.volume * self.density], dtype=np.float64) + self.volume = np.float64(np.pi * base_radius * base_radius * base_length) + self.mass = np.float64(self.volume * self.density) # Second moment of inertia area = np.pi * base_radius * base_radius diff --git a/elastica/rigidbody/mesh_rigid_body.py b/elastica/rigidbody/mesh_rigid_body.py index f83b11c54..f72925b01 100644 --- a/elastica/rigidbody/mesh_rigid_body.py +++ b/elastica/rigidbody/mesh_rigid_body.py @@ -44,7 +44,7 @@ def __init__( self.density = density self.volume = volume - self.mass = np.array([self.volume * self.density]) + self.mass = np.float64(self.volume * self.density) self.mass_second_moment_of_inertia = mass_second_moment_of_inertia.reshape( MaxDimension.value(), MaxDimension.value(), 1 ) diff --git a/elastica/rigidbody/protocol.py b/elastica/rigidbody/protocol.py index abecdecb1..22bc049d9 100644 --- a/elastica/rigidbody/protocol.py +++ b/elastica/rigidbody/protocol.py @@ -3,13 +3,13 @@ import numpy as np from numpy.typing import NDArray -from elastica.systems.protocol import SystemProtocol +from elastica.systems.protocol import SystemProtocol, SlenderBodyGeometryProtocol -class RigidBodyProtocol(SystemProtocol, Protocol): +class RigidBodyProtocol(SystemProtocol, SlenderBodyGeometryProtocol, Protocol): - mass: NDArray[np.floating] - volume: NDArray[np.floating] + mass: np.floating + volume: np.floating length: np.floating tangents: NDArray[np.floating] radius: np.floating diff --git a/elastica/rigidbody/rigid_body.py b/elastica/rigidbody/rigid_body.py index 7ce554730..900d9978d 100644 --- a/elastica/rigidbody/rigid_body.py +++ b/elastica/rigidbody/rigid_body.py @@ -34,7 +34,14 @@ def __init__(self) -> None: self.external_forces: f_arr_t self.external_torques: f_arr_t - self.mass: f_arr_t + self.internal_forces: f_arr_t + self.internal_torques: f_arr_t + + self.mass: np.floating + self.volume: np.floating + self.length: np.floating + self.tangents: f_arr_t + self.radius: np.floating self.mass_second_moment_of_inertia: f_arr_t self.inv_mass_second_moment_of_inertia: f_arr_t @@ -66,6 +73,12 @@ def zeroed_out_external_forces_and_torques(self, time: float_t) -> None: self.external_forces *= 0.0 self.external_torques *= 0.0 + def compute_internal_forces_and_torques(self, time: np.floating) -> None: + """ + For rigid body, there is no internal forces and torques + """ + pass + def compute_position_center_of_mass(self) -> f_arr_t: """ Return positional center of mass diff --git a/elastica/rigidbody/sphere.py b/elastica/rigidbody/sphere.py index 92c68ac96..675571a55 100644 --- a/elastica/rigidbody/sphere.py +++ b/elastica/rigidbody/sphere.py @@ -37,10 +37,10 @@ def __init__(self, center: f_arr_t, base_radius: float_t, density: float_t) -> N self.radius = base_radius self.density = density - self.length = 2 * base_radius + self.length = np.float64(2 * base_radius) # This is for a rigid body cylinder - self.volume = 4.0 / 3.0 * np.pi * base_radius**3 - self.mass = np.array([self.volume * self.density], dtype=np.float64) + self.volume = np.float64(4.0 / 3.0 * np.pi * base_radius**3) + self.mass = np.float64(self.volume * self.density) normal = np.array([1.0, 0.0, 0.0], dtype=np.float64).reshape(dim, 1) tangents = np.array([0.0, 0.0, 1.0], dtype=np.float64).reshape(dim, 1) binormal = _batch_cross(tangents, normal) diff --git a/elastica/rod/data_structures.py b/elastica/rod/data_structures.py index 5db3d03df..46588785b 100644 --- a/elastica/rod/data_structures.py +++ b/elastica/rod/data_structures.py @@ -9,9 +9,9 @@ from elastica._linalg import _batch_matmul if TYPE_CHECKING: - from elastica.memory_block.protocol import BlockCosseratRodProtocol + from elastica.systems.protocol import SymplecticSystemProtocol else: - BlockCosseratRodProtocol = "BlockCosseratRodProtocol" + SymplecticSystemProtocol = "SymplecticSystemProtocol" # FIXME : Explicit Stepper doesn't work as States lose the # views they initially had when working with a timestepper. @@ -50,7 +50,7 @@ class _RodSymplecticStepperMixin: - def __init__(self: BlockCosseratRodProtocol) -> None: + def __init__(self: SymplecticSystemProtocol) -> None: self.kinematic_states = _KinematicState( self.position_collection, self.director_collection ) @@ -68,7 +68,7 @@ def __init__(self: BlockCosseratRodProtocol) -> None: self.kinematic_rates = self.dynamic_states.kinematic_rates def dynamic_rates( - self: BlockCosseratRodProtocol, + self: SymplecticSystemProtocol, time: np.floating, prefac: np.floating, ) -> NDArray[np.floating]: diff --git a/elastica/rod/protocol.py b/elastica/rod/protocol.py index e1a6c0552..62d33fe57 100644 --- a/elastica/rod/protocol.py +++ b/elastica/rod/protocol.py @@ -3,7 +3,7 @@ import numpy as np from numpy.typing import NDArray -from elastica.systems.protocol import SystemProtocol +from elastica.systems.protocol import SystemProtocol, SlenderBodyGeometryProtocol class _RodEnergy(Protocol): @@ -12,7 +12,9 @@ def compute_bending_energy(self) -> NDArray[np.float64]: ... def compute_shear_energy(self) -> NDArray[np.float64]: ... -class CosseratRodProtocol(SystemProtocol, _RodEnergy, Protocol): +class CosseratRodProtocol( + SystemProtocol, SlenderBodyGeometryProtocol, _RodEnergy, Protocol +): mass: NDArray[np.floating] volume: NDArray[np.floating] @@ -44,3 +46,4 @@ class CosseratRodProtocol(SystemProtocol, _RodEnergy, Protocol): ring_rod_flag: bool periodic_boundary_nodes_idx: NDArray[np.integer] periodic_boundary_elems_idx: NDArray[np.integer] + periodic_boundary_voronoi_idx: NDArray[np.integer] diff --git a/elastica/rod/rod_base.py b/elastica/rod/rod_base.py index 76e262f15..11dc9211f 100644 --- a/elastica/rod/rod_base.py +++ b/elastica/rod/rod_base.py @@ -32,3 +32,7 @@ def __init__(self) -> None: self.ghost_voronoi_idx: NDArray[np.integer] self.ghost_elems_idx: NDArray[np.integer] + + self.periodic_boundary_nodes_idx: NDArray[np.integer] + self.periodic_boundary_elems_idx: NDArray[np.integer] + self.periodic_boundary_voronoi_idx: NDArray[np.integer] diff --git a/elastica/systems/analytical.py b/elastica/systems/analytical.py index d12cea614..317b3f99e 100644 --- a/elastica/systems/analytical.py +++ b/elastica/systems/analytical.py @@ -162,10 +162,10 @@ def energy(st): current_energy = energy(self._state) return current_energy, anal_energy - def update_internal_forces_and_torques(self, time): + def compute_internal_forces_and_torques(self, time): pass - def reset_external_forces_and_torques(self, time): + def zeroed_out_external_forces_and_torques(self, time): pass diff --git a/elastica/systems/protocol.py b/elastica/systems/protocol.py index d49129743..ef4cb7a6a 100644 --- a/elastica/systems/protocol.py +++ b/elastica/systems/protocol.py @@ -21,13 +21,21 @@ def compute_velocity_center_of_mass(self) -> NDArray[np.float64]: ... def compute_position_center_of_mass(self) -> NDArray[np.float64]: ... -class SystemProtocol(_SystemWithEnergy, _SystemWithCenterOfMass, Protocol): +class SystemProtocol(Protocol): """ Protocol for all elastica system """ REQUISITE_MODULES: list[Type] + def compute_internal_forces_and_torques(self, time: np.floating) -> None: ... + + def update_accelerations(self, time: np.floating) -> None: ... + + def zeroed_out_external_forces_and_torques(self, time: np.floating) -> None: ... + + +class SlenderBodyGeometryProtocol(Protocol): @property def n_nodes(self) -> int: ... @@ -35,29 +43,21 @@ def n_nodes(self) -> int: ... def n_elems(self) -> int: ... position_collection: NDArray[np.floating] - velocity_collection: NDArray[np.floating] - acceleration_collection: NDArray[np.floating] - director_collection: NDArray[np.floating] omega_collection: NDArray[np.floating] alpha_collection: NDArray[np.floating] - - internal_forces: NDArray[np.floating] - internal_torques: NDArray[np.floating] + director_collection: NDArray[np.floating] external_forces: NDArray[np.floating] external_torques: NDArray[np.floating] - def compute_internal_forces_and_torques(self, time: np.floating) -> None: ... - - def update_accelerations(self, time: np.floating) -> None: ... - - def zeroed_out_external_forces_and_torques(self, time: np.floating) -> None: ... + internal_forces: NDArray[np.floating] + internal_torques: NDArray[np.floating] -class SymplecticSystemProtocol(SystemProtocol, Protocol): +class SymplecticSystemProtocol(SystemProtocol, SlenderBodyGeometryProtocol, Protocol): """ Protocol for system with symplectic state variables """ @@ -80,7 +80,7 @@ def dynamic_rates( ) -> NDArray[np.floating]: ... -class ExplicitSystemProtocol(SystemProtocol, Protocol): +class ExplicitSystemProtocol(SystemProtocol, SlenderBodyGeometryProtocol, Protocol): # TODO: Temporarily made to handle explicit stepper. # Need to be refactored as the explicit stepper is further developed. def __call__(self, time: np.floating, dt: np.floating) -> np.floating: ... diff --git a/tests/test_modules/test_memory_block_rigid_body.py b/tests/test_modules/test_memory_block_rigid_body.py index 7eb4882e6..07b11d041 100644 --- a/tests/test_modules/test_memory_block_rigid_body.py +++ b/tests/test_modules/test_memory_block_rigid_body.py @@ -81,7 +81,7 @@ def test_memory_block_rigid_body(n_bodies): memory_block = MemoryBlockRigidBody(systems, system_idx_list) - assert memory_block.n_bodies == n_bodies + assert memory_block.n_systems == n_bodies assert memory_block.n_elems == n_bodies assert memory_block.n_nodes == n_bodies From 1a00700afbd5b01db9957b007f1c138dcc82b18d Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Mon, 17 Jun 2024 02:13:00 -0500 Subject: [PATCH 103/134] tests: use tmp_path to save test files --- tests/test_callback_functions.py | 12 ++++++------ tests/test_restart.py | 29 +++++++++++++++++++---------- 2 files changed, 25 insertions(+), 16 deletions(-) diff --git a/tests/test_callback_functions.py b/tests/test_callback_functions.py index 8c6362b36..0e5f3e21e 100644 --- a/tests/test_callback_functions.py +++ b/tests/test_callback_functions.py @@ -104,14 +104,14 @@ def test_my_call_back_base_class(self, n_elems): class TestExportCallBackClass: @pytest.mark.parametrize("method", ["0", 1, "numba", "test", "some string", None]) - def test_export_call_back_unavailable_save_methods(self, method): + def test_export_call_back_unavailable_save_methods(self, tmp_path, method): with pytest.raises(AssertionError) as excinfo: - callback = ExportCallBack(1, "rod", "tempdir", method) + callback = ExportCallBack(1, "rod", tmp_path.as_posix(), method) @pytest.mark.parametrize("method", ExportCallBack.AVAILABLE_METHOD) - def test_export_call_back_available_save_methods(self, method): + def test_export_call_back_available_save_methods(self, tmp_path, method): try: - callback = ExportCallBack(1, "rod", "tempdir", method) + callback = ExportCallBack(1, "rod", tmp_path.as_posix(), method) except Error: pytest.fail( f"Could not create callback module with available method {method}" @@ -236,7 +236,7 @@ def test_export_call_back_clear_test(self): assert os.path.exists(saved_path_name), "File is not saved." @pytest.mark.parametrize("n_elems", [2, 4, 16]) - def test_export_call_back_class_tempfile_option(self, n_elems): + def test_export_call_back_class_tempfile_option(self, tmp_path, n_elems): """ This test case is for testing ExportCallBack function, saving into temporary files. """ @@ -256,7 +256,7 @@ def test_export_call_back_class_tempfile_option(self, n_elems): } callback = ExportCallBack( - step_skip, "rod", "tempdir", "tempfile", file_save_interval=10 + step_skip, "rod", tmp_path.as_posix(), "tempfile", file_save_interval=10 ) for i in range(10): callback.make_callback(mock_rod, time[i], current_step[i]) diff --git a/tests/test_restart.py b/tests/test_restart.py index 0f814b21a..f19cf0377 100644 --- a/tests/test_restart.py +++ b/tests/test_restart.py @@ -48,13 +48,13 @@ def load_collection(self): return sc, rod_list - def test_restart_save_load(self, load_collection): + def test_restart_save_load(self, tmp_path, load_collection): simulator_class, rod_list = load_collection # Finalize simulator simulator_class.finalize() - directory = "restart_test_data/" + directory = (tmp_path / "restart_test_data").as_posix() time = np.random.rand() # save state @@ -79,7 +79,7 @@ def test_restart_save_load(self, load_collection): assert_allclose(test_value, correct_value) - def run_sim(self, final_time, load_from_restart, save_data_restart): + def run_sim(self, final_time, load_from_restart, save_data_restart, tmp_path): class BaseSimulatorClass( BaseSystemCollection, Constraints, Forcing, Connections, CallBacks ): @@ -118,7 +118,7 @@ class BaseSimulatorClass( # Finalize simulator simulator_class.finalize() - directory = "restart_test_data/" + directory = (tmp_path / "restart_test_data").as_posix() time_step = 1e-4 total_steps = int(final_time / time_step) @@ -149,22 +149,31 @@ class BaseSimulatorClass( return recorded_list @pytest.mark.parametrize("final_time", [0.2, 1.0]) - def test_save_restart_run_sim(self, final_time): + def test_save_restart_run_sim(self, tmp_path, final_time): # First half of simulation _ = self.run_sim( - final_time / 2, load_from_restart=False, save_data_restart=True + final_time / 2, + load_from_restart=False, + save_data_restart=True, + tmp_path=tmp_path, ) # Second half of simulation recorded_list = self.run_sim( - final_time / 2, load_from_restart=True, save_data_restart=False + final_time / 2, + load_from_restart=True, + save_data_restart=False, + tmp_path=tmp_path, ) recorded_list_second_half = recorded_list.copy() # Full simulation recorded_list = self.run_sim( - final_time, load_from_restart=False, save_data_restart=False + final_time, + load_from_restart=False, + save_data_restart=False, + tmp_path=tmp_path, ) recorded_list_full_sim = recorded_list.copy() @@ -197,13 +206,13 @@ def load_collection(self): return sc, cylinder_list - def test_restart_save_load(self, load_collection): + def test_restart_save_load(self, tmp_path, load_collection): simulator_class, cylinder_list = load_collection # Finalize simulator simulator_class.finalize() - directory = "restart_test_data/" + directory = (tmp_path / "restart_test_data").as_posix() time = np.random.rand() # save state From c1d4bdc31bfb45795dcdcf42a47247399f34b5d2 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Mon, 17 Jun 2024 02:20:27 -0500 Subject: [PATCH 104/134] fix: formatting --- elastica/callback_functions.py | 4 +--- elastica/collision/AABBCollection.py | 4 ++-- elastica/interaction.py | 1 - elastica/rod/knot_theory.py | 4 +--- 4 files changed, 4 insertions(+), 9 deletions(-) diff --git a/elastica/callback_functions.py b/elastica/callback_functions.py index 8cfdab9ce..8e89683de 100644 --- a/elastica/callback_functions.py +++ b/elastica/callback_functions.py @@ -1,6 +1,6 @@ __doc__ = """ Module contains callback classes to save simulation data for rod-like objects """ from typing import Any, Optional, TypeVar, Generic -from elastica.typing import SystemType +from elastica.typing import RodType, RigidBodyType, SystemType import os import sys @@ -11,8 +11,6 @@ from collections import defaultdict -from elastica.typing import RodType, RigidBodyType, SystemType - T = TypeVar("T") diff --git a/elastica/collision/AABBCollection.py b/elastica/collision/AABBCollection.py index c83fccac0..9385dab40 100644 --- a/elastica/collision/AABBCollection.py +++ b/elastica/collision/AABBCollection.py @@ -40,7 +40,7 @@ def __init__( @classmethod def make_from_aabb( - cls, aabb_collection: list[AABBCollection], scale_factor: int = 4 + cls, aabb_collection: list["AABBCollection"], scale_factor: int = 4 ) -> Self: # Make position collection and dimension collection arrays from aabb_collection # Wasted effort, but only once during construction @@ -64,7 +64,7 @@ def make_from_aabb( return cls(elemental_position_collection, dimension_collection, scale_factor) - def _update(self, aabb_collection: list[AABBCollection]) -> None: + def _update(self, aabb_collection: list["AABBCollection"]) -> None: # Updates internal state from another aabb """ diff --git a/elastica/interaction.py b/elastica/interaction.py index 784aab5c3..24ecb92e7 100644 --- a/elastica/interaction.py +++ b/elastica/interaction.py @@ -1,7 +1,6 @@ __doc__ = """ Numba implementation module containing interactions between a rod and its environment.""" -pass import numpy as np from elastica.external_forces import NoForces from numba import njit diff --git a/elastica/rod/knot_theory.py b/elastica/rod/knot_theory.py index 404131e35..ac440d140 100644 --- a/elastica/rod/knot_theory.py +++ b/elastica/rod/knot_theory.py @@ -10,11 +10,9 @@ The details discussion is included in `N Charles et. al. PRL (2019) `_. """ -pass - -from numba import njit import numpy as np from numpy.typing import NDArray +from numba import njit from elastica.rod.rod_base import RodBase from elastica._linalg import _batch_norm, _batch_dot, _batch_cross From c323e56a1cd3cfae5f6428f4f15b33b8c79dcf7e Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Mon, 17 Jun 2024 07:51:49 -0500 Subject: [PATCH 105/134] fix: directly access number of elements --- elastica/modules/connections.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/elastica/modules/connections.py b/elastica/modules/connections.py index d90173853..71b15834d 100644 --- a/elastica/modules/connections.py +++ b/elastica/modules/connections.py @@ -68,8 +68,8 @@ def connect( # For each system identified, get max dofs sys_idx_first = self._get_sys_idx_if_valid(first_rod) sys_idx_second = self._get_sys_idx_if_valid(second_rod) - sys_dofs_first = self._systems[sys_idx_first].n_elems - sys_dofs_second = self._systems[sys_idx_second].n_elems + sys_dofs_first = first_rod.n_elems + sys_dofs_second = second_rod.n_elems # Create _Connect object, cache it and return to user _connect: ModuleProtocol = _Connect( From 6c20c3a79fa6314bf7831435b96715f594321ecf Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Mon, 17 Jun 2024 08:31:50 -0500 Subject: [PATCH 106/134] remove indexing feature for connection --- tests/test_modules/test_connections.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/test_modules/test_connections.py b/tests/test_modules/test_connections.py index 334ece41b..53ed22f46 100644 --- a/tests/test_modules/test_connections.py +++ b/tests/test_modules/test_connections.py @@ -313,15 +313,15 @@ def mock_init(self, *args, **kwargs): ) # Constrain any and all systems - system_collection_with_connections.connect(0, 1).using( - MockConnect, 2, 42 - ) # index based connect + # system_collection_with_connections.connect(0, 1).using( + # MockConnect, 2, 42 + # ) # index based connect system_collection_with_connections.connect(mock_rod_one, mock_rod_two).using( MockConnect, 2, 3 ) # system based connect - system_collection_with_connections.connect(0, mock_rod_one).using( - MockConnect, 1, 2 - ) # index/system based connect + # system_collection_with_connections.connect(0, mock_rod_one).using( + # MockConnect, 1, 2 + # ) # index/system based connect return system_collection_with_connections, MockConnect From d32421794e266869e7ebd343fc269d89163342b9 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Mon, 17 Jun 2024 09:39:39 -0500 Subject: [PATCH 107/134] refactor: rigid body numbers typing --- elastica/rigidbody/__init__.py | 2 -- elastica/rigidbody/_typing.py | 7 ------ elastica/rigidbody/cylinder.py | 27 ++++++++++---------- elastica/rigidbody/rigid_body.py | 43 ++++++++++++++++---------------- elastica/rigidbody/sphere.py | 18 ++++++++----- 5 files changed, 47 insertions(+), 50 deletions(-) delete mode 100644 elastica/rigidbody/_typing.py diff --git a/elastica/rigidbody/__init__.py b/elastica/rigidbody/__init__.py index 9d3d4f979..55c6b9f69 100644 --- a/elastica/rigidbody/__init__.py +++ b/elastica/rigidbody/__init__.py @@ -2,5 +2,3 @@ from .cylinder import Cylinder from .sphere import Sphere from .mesh_rigid_body import MeshRigidBody - -from ._typing import float_t, f_arr_t, int_t diff --git a/elastica/rigidbody/_typing.py b/elastica/rigidbody/_typing.py deleted file mode 100644 index 941e1d411..000000000 --- a/elastica/rigidbody/_typing.py +++ /dev/null @@ -1,7 +0,0 @@ -import numpy as np -from numpy.typing import NDArray -from typing import TypeAlias - -float_t: TypeAlias = np.floating -f_arr_t: TypeAlias = NDArray[float_t] -int_t: TypeAlias = int diff --git a/elastica/rigidbody/cylinder.py b/elastica/rigidbody/cylinder.py index 053dfdee3..b9ff1e0a3 100644 --- a/elastica/rigidbody/cylinder.py +++ b/elastica/rigidbody/cylinder.py @@ -1,23 +1,24 @@ -__doc__ = """""" +__doc__ = """ +Implementation of a rigid body cylinder. +""" import numpy as np +from numpy.typing import NDArray from elastica._linalg import _batch_cross from elastica.utils import MaxDimension from elastica.rigidbody.rigid_body import RigidBodyBase -from ._typing import f_arr_t, float_t, int_t - class Cylinder(RigidBodyBase): def __init__( self, - start: f_arr_t, - direction: f_arr_t, - normal: f_arr_t, - base_length: float_t, - base_radius: float_t, - density: float_t, + start: NDArray[np.floating], + direction: NDArray[np.floating], + normal: NDArray[np.floating], + base_length: np.floating, + base_radius: np.floating, + density: np.floating, ) -> None: """ Rigid body cylinder initializer. @@ -33,7 +34,7 @@ def __init__( """ def _check_array_size( - to_check: f_arr_t, name: str, expected: int_t = 3 + to_check: NDArray[np.floating], name: str, expected: int = 3 ) -> None: array_size = to_check.size assert array_size == expected, ( @@ -42,7 +43,7 @@ def _check_array_size( ) def _check_lower_bound( - to_check: float_t, name: str, lower_bound: float_t = np.float64(0.0) + to_check: np.floating, name: str, lower_bound: np.floating = np.float64(0.0) ) -> None: assert ( to_check > lower_bound @@ -60,7 +61,7 @@ def _check_lower_bound( # rigid body does not have elements it only has one node. We are setting n_elems to # zero for only make code to work. _bootstrap_from_data requires n_elems to be define - self.n_elem: int_t = 1 + self.n_elem: int = 1 normal = normal.reshape((3, 1)) tangents = direction.reshape((3, 1)) @@ -69,7 +70,7 @@ def _check_lower_bound( self.length = base_length self.density = density - dim: int_t = MaxDimension.value() + dim: int = MaxDimension.value() # This is for a rigid body cylinder self.volume = np.float64(np.pi * base_radius * base_radius * base_length) diff --git a/elastica/rigidbody/rigid_body.py b/elastica/rigidbody/rigid_body.py index 900d9978d..c6c478151 100644 --- a/elastica/rigidbody/rigid_body.py +++ b/elastica/rigidbody/rigid_body.py @@ -2,12 +2,11 @@ from typing import Type -import numpy as np from abc import ABC -from elastica._linalg import _batch_matvec, _batch_cross -from ._typing import f_arr_t, float_t -from typing import Any +import numpy as np +from numpy.typing import NDArray +from elastica._linalg import _batch_matvec, _batch_cross class RigidBodyBase(ABC): @@ -24,29 +23,29 @@ class RigidBodyBase(ABC): def __init__(self) -> None: - self.position_collection: f_arr_t - self.velocity_collection: f_arr_t - self.acceleration_collection: f_arr_t - self.omega_collection: f_arr_t - self.alpha_collection: f_arr_t - self.director_collection: f_arr_t + self.position_collection: NDArray[np.floating] + self.velocity_collection: NDArray[np.floating] + self.acceleration_collection: NDArray[np.floating] + self.omega_collection: NDArray[np.floating] + self.alpha_collection: NDArray[np.floating] + self.director_collection: NDArray[np.floating] - self.external_forces: f_arr_t - self.external_torques: f_arr_t + self.external_forces: NDArray[np.floating] + self.external_torques: NDArray[np.floating] - self.internal_forces: f_arr_t - self.internal_torques: f_arr_t + self.internal_forces: NDArray[np.floating] + self.internal_torques: NDArray[np.floating] self.mass: np.floating self.volume: np.floating self.length: np.floating - self.tangents: f_arr_t + self.tangents: NDArray[np.floating] self.radius: np.floating - self.mass_second_moment_of_inertia: f_arr_t - self.inv_mass_second_moment_of_inertia: f_arr_t + self.mass_second_moment_of_inertia: NDArray[np.floating] + self.inv_mass_second_moment_of_inertia: NDArray[np.floating] - def update_accelerations(self, time: float_t) -> None: + def update_accelerations(self, time: np.floating) -> None: np.copyto( self.acceleration_collection, (self.external_forces) / self.mass, @@ -68,7 +67,7 @@ def update_accelerations(self, time: float_t) -> None: ), ) - def zeroed_out_external_forces_and_torques(self, time: float_t) -> None: + def zeroed_out_external_forces_and_torques(self, time: np.floating) -> None: # Reset forces and torques self.external_forces *= 0.0 self.external_torques *= 0.0 @@ -79,13 +78,13 @@ def compute_internal_forces_and_torques(self, time: np.floating) -> None: """ pass - def compute_position_center_of_mass(self) -> f_arr_t: + def compute_position_center_of_mass(self) -> NDArray[np.floating]: """ Return positional center of mass """ return self.position_collection[..., 0].copy() - def compute_translational_energy(self) -> Any: + def compute_translational_energy(self) -> NDArray[np.floating]: """ Return translational energy """ @@ -97,7 +96,7 @@ def compute_translational_energy(self) -> Any: ) ) - def compute_rotational_energy(self) -> Any: + def compute_rotational_energy(self) -> NDArray[np.floating]: """ Return rotational energy """ diff --git a/elastica/rigidbody/sphere.py b/elastica/rigidbody/sphere.py index 675571a55..7c0e7d803 100644 --- a/elastica/rigidbody/sphere.py +++ b/elastica/rigidbody/sphere.py @@ -1,16 +1,22 @@ -__doc__ = """""" +__doc__ = """ +Implementation of a sphere rigid body. +""" import numpy as np +from numpy.typing import NDArray from elastica._linalg import _batch_cross from elastica.utils import MaxDimension from elastica.rigidbody.rigid_body import RigidBodyBase -from ._typing import float_t, f_arr_t, int_t - class Sphere(RigidBodyBase): - def __init__(self, center: f_arr_t, base_radius: float_t, density: float_t) -> None: + def __init__( + self, + center: NDArray[np.floating], + base_radius: np.floating, + density: np.floating, + ) -> None: """ Rigid body sphere initializer. @@ -25,9 +31,9 @@ def __init__(self, center: f_arr_t, base_radius: float_t, density: float_t) -> N # rigid body does not have elements it only have one node. We are setting n_elems to # zero for only make code to work. _bootstrap_from_data requires n_elems to be defined - self.n_elems: int_t = 1 + self.n_elems: int = 1 - dim: int_t = MaxDimension.value() + dim: int = MaxDimension.value() assert ( center.size == dim From 9267f38a0def9a3c77010670fc2fa22055b5e43a Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Mon, 17 Jun 2024 12:31:04 -0500 Subject: [PATCH 108/134] add static system type --- elastica/memory_block/protocol.py | 8 -------- elastica/modules/base_system.py | 8 +++++--- elastica/modules/constraints.py | 13 ++++++++++--- elastica/modules/contact.py | 5 +++-- elastica/modules/memory_block.py | 14 ++++++++++---- elastica/modules/protocol.py | 8 +++++--- elastica/surface/surface_base.py | 8 +++++++- elastica/systems/protocol.py | 18 ++++-------------- elastica/typing.py | 27 ++++++--------------------- 9 files changed, 50 insertions(+), 59 deletions(-) diff --git a/elastica/memory_block/protocol.py b/elastica/memory_block/protocol.py index 279f01580..5da6331b4 100644 --- a/elastica/memory_block/protocol.py +++ b/elastica/memory_block/protocol.py @@ -12,11 +12,3 @@ class BlockSystemProtocol(SystemProtocol, Protocol): @property def n_systems(self) -> int: """Number of systems in the block.""" - - -class BlockCosseratRodProtocol(CosseratRodProtocol, SymplecticSystemProtocol, Protocol): - pass - - -class BlockRigidBodyProtocol(RigidBodyProtocol, SymplecticSystemProtocol, Protocol): - pass diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index c31581d66..ff32cf922 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -8,12 +8,12 @@ from typing import Type, Generator, Iterable, Any, overload from typing import final from elastica.typing import ( + StaticSystemType, SystemType, SystemIdxType, OperatorType, OperatorCallbackType, OperatorFinalizeType, - BlockType, ) import numpy as np @@ -68,7 +68,7 @@ def __init__(self) -> None: ) # List of systems to be integrated - self._systems: list[SystemType] = [] + self._systems: list[StaticSystemType] = [] self.__final_systems: list[SystemType] = [] # Flag Finalize: Finalizing twice will cause an error, @@ -138,7 +138,9 @@ def override_allowed_types( self.allowed_sys_types = allowed_types @final - def _get_sys_idx_if_valid(self, sys_to_be_added: SystemType) -> SystemIdxType: + def _get_sys_idx_if_valid( + self, sys_to_be_added: "SystemType | StaticSystemType" + ) -> SystemIdxType: n_systems = len(self) # Total number of systems from mixed-in class sys_idx: SystemIdxType diff --git a/elastica/modules/constraints.py b/elastica/modules/constraints.py index bfa452124..344306b31 100644 --- a/elastica/modules/constraints.py +++ b/elastica/modules/constraints.py @@ -4,14 +4,20 @@ Provides the constraints interface to enforce displacement boundary conditions (see `boundary_conditions.py`). """ -from typing import Any, Type +from typing import Any, Type, cast from typing_extensions import Self import numpy as np from elastica.boundary_conditions import ConstraintBase -from elastica.typing import SystemIdxType, ConstrainingIndex, RigidBodyType, RodType +from elastica.typing import ( + SystemIdxType, + ConstrainingIndex, + RigidBodyType, + RodType, + BlockSystemType, +) from .protocol import SystemCollectionProtocol, ModuleProtocol @@ -73,7 +79,8 @@ def _finalize_constraints(self: SystemCollectionProtocol) -> None: # Apply the constrain to synchronize the periodic boundaries of the memory rod. Find the memory block # sys idx among other systems added and then apply boundary conditions. memory_block_idx = self._get_sys_idx_if_valid(block) - self.constrain(self._systems[memory_block_idx]).using( + block_system = cast(BlockSystemType, self._systems[memory_block_idx]) + self.constrain(block_system).using( _ConstrainPeriodicBoundaries, ) diff --git a/elastica/modules/contact.py b/elastica/modules/contact.py index 69cb11db0..f52bdc31e 100644 --- a/elastica/modules/contact.py +++ b/elastica/modules/contact.py @@ -10,6 +10,7 @@ from elastica.typing import ( SystemIdxType, OperatorFinalizeType, + StaticSystemType, SystemType, ) from .protocol import SystemCollectionProtocol, ModuleProtocol @@ -47,7 +48,7 @@ def __init__(self: SystemCollectionProtocol) -> None: def detect_contact_between( self: SystemCollectionProtocol, first_system: SystemType, - second_system: SystemType, + second_system: "SystemType | StaticSystemType", ) -> ModuleProtocol: """ This method adds contact detection between two objects using the selected contact class. @@ -56,7 +57,7 @@ def detect_contact_between( Parameters ---------- first_system : SystemType - second_system : SystemType + second_system : SystemType | StaticSystemType Returns ------- diff --git a/elastica/modules/memory_block.py b/elastica/modules/memory_block.py index 3f69b57e0..d0e751c15 100644 --- a/elastica/modules/memory_block.py +++ b/elastica/modules/memory_block.py @@ -6,6 +6,8 @@ from elastica.typing import ( RodType, RigidBodyType, + SurfaceType, + StaticSystemType, SystemType, SystemIdxType, BlockSystemType, @@ -18,7 +20,9 @@ from elastica.memory_block.memory_block_rigid_body import MemoryBlockRigidBody -def construct_memory_block_structures(systems: list[SystemType]) -> list[SystemType]: +def construct_memory_block_structures( + systems: list[StaticSystemType], +) -> list[SystemType]: """ This function takes the systems (rod or rigid body) appended to the simulator class and separates them into lists depending on if system is Cosserat rod or rigid body. Then using @@ -47,9 +51,11 @@ def construct_memory_block_structures(systems: list[SystemType]) -> list[SystemT temp_list_for_rigid_body_systems_idx.append(system_idx) elif isinstance(sys_to_be_added, SurfaceBase): - raise NotImplementedError( - "Surfaces are not yet implemented in memory block construction." - ) + pass + # surface_system = cast(SurfaceType, sys_to_be_added) + # raise NotImplementedError( + # "Surfaces are not yet implemented in memory block construction." + # ) else: raise TypeError( diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index b65419158..5de97c7e2 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -8,9 +8,9 @@ OperatorType, OperatorCallbackType, OperatorFinalizeType, + StaticSystemType, SystemType, ConnectionIndex, - BlockType, ) from elastica.joint import FreeJoint from elastica.callback_functions import CallBackBaseClass @@ -34,7 +34,7 @@ def id(self) -> Any: ... class SystemCollectionProtocol(Protocol): - _systems: list[SystemType] + _systems: list[StaticSystemType] def __len__(self) -> int: ... @@ -69,7 +69,9 @@ def _feature_group_finalize(self) -> list[OperatorFinalizeType]: ... def systems(self) -> Generator[SystemType, None, None]: ... - def _get_sys_idx_if_valid(self, sys_to_be_added: SystemType) -> SystemIdxType: ... + def _get_sys_idx_if_valid( + self, sys_to_be_added: "SystemType | StaticSystemType" + ) -> SystemIdxType: ... # Connection API _finalize_connections: OperatorFinalizeType diff --git a/elastica/surface/surface_base.py b/elastica/surface/surface_base.py index c5c7a2ff7..9870afe96 100644 --- a/elastica/surface/surface_base.py +++ b/elastica/surface/surface_base.py @@ -1,5 +1,5 @@ __doc__ = """Base class for surfaces""" -from typing import Type +from typing import TYPE_CHECKING, Type import numpy as np from numpy.typing import NDArray @@ -23,3 +23,9 @@ def __init__(self) -> None: """ self.normal: NDArray[np.floating] # (3,) self.origin: NDArray[np.floating] # (3, 1) + + +if TYPE_CHECKING: + from elastica.systems.protocol import StaticSystemProtocol + + _: StaticSystemProtocol = SurfaceBase() diff --git a/elastica/systems/protocol.py b/elastica/systems/protocol.py index ef4cb7a6a..ccac73f64 100644 --- a/elastica/systems/protocol.py +++ b/elastica/systems/protocol.py @@ -9,25 +9,15 @@ from numpy.typing import NDArray -class _SystemWithEnergy(Protocol): - def compute_translational_energy(self) -> NDArray[np.float64]: ... - - def compute_rotational_energy(self) -> NDArray[np.float64]: ... - - -class _SystemWithCenterOfMass(Protocol): - def compute_velocity_center_of_mass(self) -> NDArray[np.float64]: ... - - def compute_position_center_of_mass(self) -> NDArray[np.float64]: ... +class StaticSystemProtocol(Protocol): + REQUISITE_MODULES: list[Type] -class SystemProtocol(Protocol): +class SystemProtocol(StaticSystemProtocol, Protocol): """ - Protocol for all elastica system + Protocol for all dynamic elastica system """ - REQUISITE_MODULES: list[Type] - def compute_internal_forces_and_torques(self, time: np.floating) -> None: ... def update_accelerations(self, time: np.floating) -> None: ... diff --git a/elastica/typing.py b/elastica/typing.py index d13e927dc..aa9928daa 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -4,8 +4,7 @@ """ from typing import TYPE_CHECKING -from typing import Callable, Any, ParamSpec -from typing import TypeAlias +from typing import Callable, Any, ParamSpec, TypeAlias import numpy as np @@ -22,6 +21,7 @@ from .rod.data_structures import _State as State from .systems.protocol import ( SystemProtocol, + StaticSystemProtocol, SymplecticSystemProtocol, ExplicitSystemProtocol, ) @@ -30,32 +30,18 @@ SymplecticStepperProtocol, MemoryProtocol, ) - from .memory_block.protocol import ( - BlockCosseratRodProtocol, - BlockRigidBodyProtocol, - BlockSystemProtocol, - ) - - # Modules Base Classes - from .boundary_conditions import FreeBC - from .callback_functions import CallBackBaseClass - from .contact_forces import NoContact - from .dissipation import DamperBase - from .external_forces import NoForces - from .joint import FreeJoint + from .memory_block.protocol import BlockSystemProtocol from .mesh.protocol import MeshProtocol +StaticSystemType: TypeAlias = "StaticSystemProtocol" SystemType: TypeAlias = "SystemProtocol" -BlockSystemType: TypeAlias = "BlockSystemProtocol" SystemIdxType: TypeAlias = int +BlockSystemType: TypeAlias = "BlockSystemProtocol" -# ModuleObjectTypes: TypeAlias = ( -# NoForces | NoContact | FreeJoint | FreeBC | DamperBase | CallBackBaseClass -# ) +# NoForces | NoContact | FreeJoint | FreeBC | DamperBase | CallBackBaseClass -# TODO: Modify this line and move to elastica/typing.py once system state is defined # Mostly used in explicit stepper: for symplectic, use kinetic and dynamic state StateType: TypeAlias = "State" @@ -69,7 +55,6 @@ SurfaceType: TypeAlias = "SurfaceBase" SystemCollectionType: TypeAlias = "SystemCollectionProtocol" -BlockType: TypeAlias = "BlockCosseratRodProtocol | BlockRigidBodyProtocol" # Indexing types # TODO: Maybe just use slice?? From 13fd73767b162dd3424f4e6d90fe6581a9466470 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Wed, 19 Jun 2024 14:00:28 -0500 Subject: [PATCH 109/134] exclude pyvista untyped import --- elastica/mesh/mesh_initializer.py | 21 +++++++++++---------- pyproject.toml | 5 ++++- 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/elastica/mesh/mesh_initializer.py b/elastica/mesh/mesh_initializer.py index 2320c1858..b0cf32aa4 100644 --- a/elastica/mesh/mesh_initializer.py +++ b/elastica/mesh/mesh_initializer.py @@ -1,7 +1,8 @@ __doc__ = """ Mesh Initializer using Pyvista """ -import pyvista as pv +import pyvista as pv # type: ignore import numpy as np +from numpy.typing import NDArray class Mesh: @@ -100,8 +101,8 @@ def mesh_update(self) -> None: ) def face_calculation( - self, pvfaces: np.ndarray, meshpoints: np.ndarray, n_faces: int - ) -> np.ndarray: + self, pvfaces: NDArray, meshpoints: NDArray, n_faces: int + ) -> NDArray: """ This function converts the faces from pyvista to pyelastica geometry @@ -142,7 +143,7 @@ def face_calculation( return faces - def face_normal_calculation(self, pyvista_face_normals: np.ndarray) -> np.ndarray: + def face_normal_calculation(self, pyvista_face_normals: NDArray) -> NDArray: """ This function converts the face normals from pyvista to pyelastica geometry, in pyelastica the face are stored in the format of (n_faces, 3 spatial coordinates), @@ -152,7 +153,7 @@ def face_normal_calculation(self, pyvista_face_normals: np.ndarray) -> np.ndarra return face_normals - def face_center_calculation(self, faces: np.ndarray, n_faces: int) -> np.ndarray: + def face_center_calculation(self, faces: NDArray, n_faces: int) -> NDArray: """ This function calculates the position vector of each face of the mesh simply by averaging all the vertices of every face/cell. @@ -167,7 +168,7 @@ def face_center_calculation(self, faces: np.ndarray, n_faces: int) -> np.ndarray return face_centers - def mesh_scale_calculation(self, bounds: np.ndarray) -> np.ndarray: + def mesh_scale_calculation(self, bounds: NDArray) -> NDArray: """ This function calculates scale of the mesh, for that it calculates the maximum distance between mesh's farthest verticies in each axis. @@ -180,7 +181,7 @@ def mesh_scale_calculation(self, bounds: np.ndarray) -> np.ndarray: return scale - def orientation_calculation(self, cube_face_normals: np.ndarray) -> np.ndarray: + def orientation_calculation(self, cube_face_normals: NDArray) -> NDArray: """ This function calculates the orientation of the mesh by using a dummy cube utility from pyvista. @@ -209,7 +210,7 @@ def visualize(self) -> None: pyvista_plotter.add_mesh(self.mesh) pyvista_plotter.show() - def translate(self, target_center: np.ndarray) -> None: + def translate(self, target_center: NDArray) -> None: """ Parameters: {numpy.ndarray-(3 spatial coordinates)} ex : mesh.translate(np.array([1,1,1])) @@ -220,7 +221,7 @@ def translate(self, target_center: np.ndarray) -> None: self.mesh = self.mesh.translate(target_center) self.mesh_update() - def scale(self, factor: np.ndarray) -> None: + def scale(self, factor: NDArray) -> None: """ Parameters: {numpy.ndarray-(3 spatial constants)} ex : mesh.scale(np.array([1,1,1])) @@ -230,7 +231,7 @@ def scale(self, factor: np.ndarray) -> None: self.mesh = self.mesh.scale(factor) self.mesh_update() - def rotate(self, axis: np.ndarray, angle: float) -> None: + def rotate(self, axis: NDArray, angle: float) -> None: """ Parameters: {rotation_axis: unit vector[numpy.ndarray-(3 spatial coordinates)], angle: in degrees[float]} ex : mesh.rotate(np.array([1,0,0]), 90) diff --git a/pyproject.toml b/pyproject.toml index 2628c4f9a..9851cafc9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -145,7 +145,10 @@ warn_unused_ignores = false exclude = [ "elastica/systems/analytical.py", "elastica/experimental/*", - "elastica/mesh/mesh_initializer.py", +] + +untyped_calls_exclude = [ + "pyvista", ] [tool.coverage.report] From e54536c43bce3bb5a2643803b175b0e2877f81cd Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Mon, 24 Jun 2024 09:36:57 -0500 Subject: [PATCH 110/134] docs: add duck-typing structure into documentation --- docs/advanced/PackageDesign.md | 110 ++++++++++++++++++++++++++++++++- docs/conf.py | 4 ++ poetry.lock | 30 ++++----- pyproject.toml | 2 + 4 files changed, 129 insertions(+), 17 deletions(-) diff --git a/docs/advanced/PackageDesign.md b/docs/advanced/PackageDesign.md index 857a3c046..1ef461214 100644 --- a/docs/advanced/PackageDesign.md +++ b/docs/advanced/PackageDesign.md @@ -1,8 +1,114 @@ -# Code Design: Mixin and Composition +# Code Design + +## Mixin and Composition Elastica package follows Mixin and composition design patterns that may be unfamiliar to users. Here is a collection of references that introduce the package design. -## References +### References - [stackoverflow discussion on Mixin](https://stackoverflow.com/questions/533631/what-is-a-mixin-and-why-are-they-useful) - [example of Mixin: python collections](https://docs.python.org/dev/library/collections.abc.html) + +## Duck Typing + +Elastica package uses duck typing to allow users to define their own classes and functions. Here is a `typing.Protocol` structure that is used in the package. + +### Systems + +``` {mermaid} + flowchart LR + direction RL + subgraph Systems Protocol + direction RL + SLBD(SlenderBodyGeometryProtool) + SymST["SymplecticSystem:\n• KinematicStates/Rates\n• DynamicStates/Rates"] + style SymST text-align:left + ExpST["ExplicitSystem:\n• States (Unused)"] + style ExpST text-align:left + P((position\nvelocity\nacceleration\n..)) --> SLBD + subgraph StaticSystemType + Surface + Mesh + end + subgraph SystemType + direction TB + Rod + RigidBody + end + SLBD --> SymST + SystemType --> SymST + SLBD --> ExpST + SystemType --> ExpST + end + subgraph Timestepper Protocol + direction TB + StP["StepperProtocol\n• step(SystemCollection, time, dt)"] + style StP text-align:left + SymplecticStepperProtocol["SymplecticStepperProtocol\n• PositionVerlet"] + style SymplecticStepperProtocol text-align:left + ExpplicitStepperProtocol["ExpplicitStepperProtocol\n(Unused)"] + end + + subgraph SystemCollection + + end + SymST --> SystemCollection --> SymplecticStepperProtocol + ExpST --> SystemCollection --> ExpplicitStepperProtocol + StaticSystemType --> SystemCollection + +``` + +### System Collection (Build memory block) + +``` {mermaid} + flowchart LR + Sys((Systems)) + St((Stepper)) + subgraph SystemCollectionType + direction LR + StSys["StaticSystem:\n• Surface\n• Mesh"] + style StSys text-align:left + DynSys["DynamicSystem:\n• Rod\n  • CosseratRod\n• RigidBody\n  • Sphere\n  • Cylinder"] + style DynSys text-align:left + + BlDynSys["BlockSystemType:\n• BlockCosseratRod\n• BlockRigidBody"] + style BlDynSys text-align:left + + F{{"Feature Group (OperatorGroup):\n• Synchronize\n• Constrain values\n• Constrain rates\n• Callback"}} + style F text-align:left + end + Sys --> StSys --> F + Sys --> DynSys -->|Finalize| BlDynSys --> St + DynSys --> F <--> St + +``` + +### System Collection (Features) + +``` {mermaid} + flowchart LR + Sys((Systems)) + St((Stepper)) + subgraph SystemCollectionType + direction LR + StSys["StaticSystem:\n• Surface\n• Mesh"] + style StSys text-align:left + DynSys["DynamicSystem:\n• Rod\n  • CosseratRod\n• RigidBody\n  • Sphere\n  • Cylinder"] + style DynSys text-align:left + + subgraph Feature + direction LR + Forcing -->|add_forcing_to| Synchronize + Constraints -->|constrain| ConstrainValues + Constraints -->|constrain| ConstrainRates + Contact -->|detect_contact_between| Synchronize + Connection -->|connect| Synchronize + Damping -->|dampen| ConstrainRates + Callback -->|collect_diagnosis| CallbackGroup + end + end + Sys --> StSys --> Feature + Sys --> DynSys + DynSys --> Feature <--> St + +``` diff --git a/docs/conf.py b/docs/conf.py index 6d2047961..f2ca1f397 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -41,6 +41,7 @@ #'sphinx.ext.napoleon', 'sphinx.ext.viewcode', 'sphinx.ext.mathjax', + "sphinxcontrib.mermaid", 'numpydoc', 'myst_parser', ] @@ -98,3 +99,6 @@ # -- Options for numpydoc --------------------------------------------------- numpydoc_show_class_members = False + +# -- Mermaid configuration --------------------------------------------------- +mermaid_params = ['--theme', 'neutral'] diff --git a/poetry.lock b/poetry.lock index ee1ffa8ba..196c8dd33 100644 --- a/poetry.lock +++ b/poetry.lock @@ -627,7 +627,7 @@ files = [ name = "jinja2" version = "3.1.4" description = "A very fast and expressive template engine." -optional = false +optional = true python-versions = ">=3.7" files = [ {file = "jinja2-3.1.4-py3-none-any.whl", hash = "sha256:bc5dd2abb727a5319567b7a813e6a2e7318c39f4f487cfe6c89c6f9c7d25197d"}, @@ -814,7 +814,7 @@ testing = ["coverage", "pytest", "pytest-cov", "pytest-regressions"] name = "markupsafe" version = "2.1.5" description = "Safely add untrusted strings to HTML/XML markup." -optional = false +optional = true python-versions = ">=3.7" files = [ {file = "MarkupSafe-2.1.5-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:a17a92de5231666cfbe003f0e4b9b3a7ae3afb1ec2845aadc2bacc93ff85febc"}, @@ -1557,7 +1557,6 @@ files = [ {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:69b023b2b4daa7548bcfbd4aa3da05b3a74b772db9e23b982788168117739938"}, {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:81e0b275a9ecc9c0c0c07b4b90ba548307583c125f54d5b6946cfee6360c733d"}, {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba336e390cd8e4d1739f42dfe9bb83a3cc2e80f567d8805e11b46f4a943f5515"}, - {file = "PyYAML-6.0.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:326c013efe8048858a6d312ddd31d56e468118ad4cdeda36c719bf5bb6192290"}, {file = "PyYAML-6.0.1-cp310-cp310-win32.whl", hash = "sha256:bd4af7373a854424dabd882decdc5579653d7868b8fb26dc7d0e99f823aa5924"}, {file = "PyYAML-6.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:fd1592b3fdf65fff2ad0004b5e363300ef59ced41c2e6b3a99d4089fa8c5435d"}, {file = "PyYAML-6.0.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:6965a7bc3cf88e5a1c3bd2e0b5c22f8d677dc88a455344035f03399034eb3007"}, @@ -1565,16 +1564,8 @@ files = [ {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:42f8152b8dbc4fe7d96729ec2b99c7097d656dc1213a3229ca5383f973a5ed6d"}, {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:062582fca9fabdd2c8b54a3ef1c978d786e0f6b3a1510e0ac93ef59e0ddae2bc"}, {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d2b04aac4d386b172d5b9692e2d2da8de7bfb6c387fa4f801fbf6fb2e6ba4673"}, - {file = "PyYAML-6.0.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e7d73685e87afe9f3b36c799222440d6cf362062f78be1013661b00c5c6f678b"}, {file = "PyYAML-6.0.1-cp311-cp311-win32.whl", hash = "sha256:1635fd110e8d85d55237ab316b5b011de701ea0f29d07611174a1b42f1444741"}, {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"}, - {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"}, - {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"}, - {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"}, - {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"}, - {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"}, - {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"}, - {file = "PyYAML-6.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:0d3304d8c0adc42be59c5f8a4d9e3d7379e6955ad754aa9d6ab7a398b59dd1df"}, {file = "PyYAML-6.0.1-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:50550eb667afee136e9a77d6dc71ae76a44df8b3e51e41b77f6de2932bfe0f47"}, {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1fe35611261b29bd1de0070f0b2f47cb6ff71fa6595c077e42bd0c419fa27b98"}, {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:704219a11b772aea0d8ecd7058d0082713c3562b4e271b849ad7dc4a5c90c13c"}, @@ -1591,7 +1582,6 @@ files = [ {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a0cd17c15d3bb3fa06978b4e8958dcdc6e0174ccea823003a106c7d4d7899ac5"}, {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:28c119d996beec18c05208a8bd78cbe4007878c6dd15091efb73a30e90539696"}, {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7e07cbde391ba96ab58e532ff4803f79c4129397514e1413a7dc761ccd755735"}, - {file = "PyYAML-6.0.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:49a183be227561de579b4a36efbb21b3eab9651dd81b1858589f796549873dd6"}, {file = "PyYAML-6.0.1-cp38-cp38-win32.whl", hash = "sha256:184c5108a2aca3c5b3d3bf9395d50893a7ab82a38004c8f61c258d4428e80206"}, {file = "PyYAML-6.0.1-cp38-cp38-win_amd64.whl", hash = "sha256:1e2722cc9fbb45d9b87631ac70924c11d3a401b2d7f410cc0e3bbf249f2dca62"}, {file = "PyYAML-6.0.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:9eb6caa9a297fc2c2fb8862bc5370d0303ddba53ba97e71f08023b6cd73d16a8"}, @@ -1599,7 +1589,6 @@ files = [ {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5773183b6446b2c99bb77e77595dd486303b4faab2b086e7b17bc6bef28865f6"}, {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b786eecbdf8499b9ca1d697215862083bd6d2a99965554781d0d8d1ad31e13a0"}, {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bc1bf2925a1ecd43da378f4db9e4f799775d6367bdb94671027b73b393a7c42c"}, - {file = "PyYAML-6.0.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:04ac92ad1925b2cff1db0cfebffb6ffc43457495c9b3c39d3fcae417d7125dc5"}, {file = "PyYAML-6.0.1-cp39-cp39-win32.whl", hash = "sha256:faca3bdcf85b2fc05d06ff3fbc1f83e1391b3e724afa3feba7d13eeab355484c"}, {file = "PyYAML-6.0.1-cp39-cp39-win_amd64.whl", hash = "sha256:510c9deebc5c0225e8c96813043e62b680ba2f9c50a08d3724c7f28a747d1486"}, {file = "PyYAML-6.0.1.tar.gz", hash = "sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43"}, @@ -1877,6 +1866,17 @@ files = [ [package.extras] test = ["flake8", "mypy", "pytest"] +[[package]] +name = "sphinxcontrib-mermaid" +version = "0.9.2" +description = "Mermaid diagrams in yours Sphinx powered docs" +optional = true +python-versions = ">=3.7" +files = [ + {file = "sphinxcontrib-mermaid-0.9.2.tar.gz", hash = "sha256:252ef13dd23164b28f16d8b0205cf184b9d8e2b714a302274d9f59eb708e77af"}, + {file = "sphinxcontrib_mermaid-0.9.2-py3-none-any.whl", hash = "sha256:6795a72037ca55e65663d2a2c1a043d636dc3d30d418e56dd6087d1459d98a5d"}, +] + [[package]] name = "sphinxcontrib-qthelp" version = "1.0.7" @@ -2044,10 +2044,10 @@ numpy = ["numpy (>=1.9)"] web = ["wslink (>=1.0.4)"] [extras] -docs = ["Sphinx", "docutils", "myst-parser", "numpydoc", "readthedocs-sphinx-search", "sphinx-autodoc-typehints", "sphinx-book-theme"] +docs = ["Sphinx", "docutils", "myst-parser", "numpydoc", "readthedocs-sphinx-search", "sphinx-autodoc-typehints", "sphinx-book-theme", "sphinxcontrib-mermaid"] examples = ["cma"] [metadata] lock-version = "2.0" python-versions = ">=3.10,<3.12" -content-hash = "07105447f85d22c61c0766e22959a4f3428f9cfa20e7499aae4017fa993fdca3" +content-hash = "c9747151a346aa7c37c83f66c82b8283aa4ed796788d545d65378b7e8a117260" diff --git a/pyproject.toml b/pyproject.toml index 9851cafc9..bbd4e9ce4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,7 @@ Sphinx = {version = "^6.1", optional = true} sphinx-book-theme = {version = "^1.0", optional = true} readthedocs-sphinx-search = {version = ">=0.1.1,<0.4.0", optional = true} sphinx-autodoc-typehints = {version = "^1.21", optional = true} +sphinxcontrib-mermaid = {version = "^0.9.2", optional = true} myst-parser = {version = "^1.0", optional = true} numpydoc = {version = "^1.3.1", optional = true} docutils = {version = "^0.18", optional = true} @@ -73,6 +74,7 @@ docs = [ "sphinx-book-theme", "readthedocs-sphinx-search", "sphinx-autodoc-typehints", + "sphinxcontrib-mermaid", "myst-parser", "numpydoc", "docutils", From 26dd70f8fb3b08817f61ca092ba46fec6dea4892 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Mon, 24 Jun 2024 09:36:57 -0500 Subject: [PATCH 111/134] docs: add duck-typing structure into documentation --- docs/advanced/PackageDesign.md | 110 ++++++++++++++++++++++++++++- docs/conf.py | 4 ++ elastica/modules/operator_group.py | 23 ++++-- poetry.lock | 30 ++++---- pyproject.toml | 2 + 5 files changed, 148 insertions(+), 21 deletions(-) diff --git a/docs/advanced/PackageDesign.md b/docs/advanced/PackageDesign.md index 857a3c046..1ef461214 100644 --- a/docs/advanced/PackageDesign.md +++ b/docs/advanced/PackageDesign.md @@ -1,8 +1,114 @@ -# Code Design: Mixin and Composition +# Code Design + +## Mixin and Composition Elastica package follows Mixin and composition design patterns that may be unfamiliar to users. Here is a collection of references that introduce the package design. -## References +### References - [stackoverflow discussion on Mixin](https://stackoverflow.com/questions/533631/what-is-a-mixin-and-why-are-they-useful) - [example of Mixin: python collections](https://docs.python.org/dev/library/collections.abc.html) + +## Duck Typing + +Elastica package uses duck typing to allow users to define their own classes and functions. Here is a `typing.Protocol` structure that is used in the package. + +### Systems + +``` {mermaid} + flowchart LR + direction RL + subgraph Systems Protocol + direction RL + SLBD(SlenderBodyGeometryProtool) + SymST["SymplecticSystem:\n• KinematicStates/Rates\n• DynamicStates/Rates"] + style SymST text-align:left + ExpST["ExplicitSystem:\n• States (Unused)"] + style ExpST text-align:left + P((position\nvelocity\nacceleration\n..)) --> SLBD + subgraph StaticSystemType + Surface + Mesh + end + subgraph SystemType + direction TB + Rod + RigidBody + end + SLBD --> SymST + SystemType --> SymST + SLBD --> ExpST + SystemType --> ExpST + end + subgraph Timestepper Protocol + direction TB + StP["StepperProtocol\n• step(SystemCollection, time, dt)"] + style StP text-align:left + SymplecticStepperProtocol["SymplecticStepperProtocol\n• PositionVerlet"] + style SymplecticStepperProtocol text-align:left + ExpplicitStepperProtocol["ExpplicitStepperProtocol\n(Unused)"] + end + + subgraph SystemCollection + + end + SymST --> SystemCollection --> SymplecticStepperProtocol + ExpST --> SystemCollection --> ExpplicitStepperProtocol + StaticSystemType --> SystemCollection + +``` + +### System Collection (Build memory block) + +``` {mermaid} + flowchart LR + Sys((Systems)) + St((Stepper)) + subgraph SystemCollectionType + direction LR + StSys["StaticSystem:\n• Surface\n• Mesh"] + style StSys text-align:left + DynSys["DynamicSystem:\n• Rod\n  • CosseratRod\n• RigidBody\n  • Sphere\n  • Cylinder"] + style DynSys text-align:left + + BlDynSys["BlockSystemType:\n• BlockCosseratRod\n• BlockRigidBody"] + style BlDynSys text-align:left + + F{{"Feature Group (OperatorGroup):\n• Synchronize\n• Constrain values\n• Constrain rates\n• Callback"}} + style F text-align:left + end + Sys --> StSys --> F + Sys --> DynSys -->|Finalize| BlDynSys --> St + DynSys --> F <--> St + +``` + +### System Collection (Features) + +``` {mermaid} + flowchart LR + Sys((Systems)) + St((Stepper)) + subgraph SystemCollectionType + direction LR + StSys["StaticSystem:\n• Surface\n• Mesh"] + style StSys text-align:left + DynSys["DynamicSystem:\n• Rod\n  • CosseratRod\n• RigidBody\n  • Sphere\n  • Cylinder"] + style DynSys text-align:left + + subgraph Feature + direction LR + Forcing -->|add_forcing_to| Synchronize + Constraints -->|constrain| ConstrainValues + Constraints -->|constrain| ConstrainRates + Contact -->|detect_contact_between| Synchronize + Connection -->|connect| Synchronize + Damping -->|dampen| ConstrainRates + Callback -->|collect_diagnosis| CallbackGroup + end + end + Sys --> StSys --> Feature + Sys --> DynSys + DynSys --> Feature <--> St + +``` diff --git a/docs/conf.py b/docs/conf.py index 6d2047961..f2ca1f397 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -41,6 +41,7 @@ #'sphinx.ext.napoleon', 'sphinx.ext.viewcode', 'sphinx.ext.mathjax', + "sphinxcontrib.mermaid", 'numpydoc', 'myst_parser', ] @@ -98,3 +99,6 @@ # -- Options for numpydoc --------------------------------------------------- numpydoc_show_class_members = False + +# -- Mermaid configuration --------------------------------------------------- +mermaid_params = ['--theme', 'neutral'] diff --git a/elastica/modules/operator_group.py b/elastica/modules/operator_group.py index 2a3975728..3b01225d4 100644 --- a/elastica/modules/operator_group.py +++ b/elastica/modules/operator_group.py @@ -11,16 +11,31 @@ class OperatorGroupFIFO(Iterable, Generic[T, F]): """ A class to store the features and their corresponding operators in a FIFO manner. + Feature can be any user-defined object to label the operators, and operators are + callable functions. + + This data structure is used to organize elastica operators, such as forcing, + constraints, boundary condition, etc. Examples -------- + >>> class FeatureObj: + ... ADD: Callable + ... SUBTRACT: Callable + ... MULTIPLY: Callable + + >>> obj_1 = FeatureObj() + >>> obj_2 = FeatureObj() + >>> >>> operator_group = OperatorGroupFIFO() >>> operator_group.append_id(obj_1) >>> operator_group.append_id(obj_2) - >>> operator_group.add_operators(obj_1, [ADD, SUBTRACT]) - >>> operator_group.add_operators(obj_2, [SUBTRACT, MULTIPLY]) - >>> list(operator_group) - [ADD, SUBTRACT, SUBTRACT, MULTIPLY] + >>> operator_group.add_operators(obj_1, [obj_1.ADD, obj_1.SUBTRACT]) + >>> operator_group.add_operators(obj_2, [obj_2.SUBTRACT, obj_2.MULTIPLY]) + >>> list(iter(operator_group)) + [obj_1.ADD, obj_1.SUBTRACT, obj_2.SUBTRACT, obj_2.MULTIPLY] + >>> for operator in operator_group: # call operator in the group + ... operator() Attributes ---------- diff --git a/poetry.lock b/poetry.lock index ee1ffa8ba..196c8dd33 100644 --- a/poetry.lock +++ b/poetry.lock @@ -627,7 +627,7 @@ files = [ name = "jinja2" version = "3.1.4" description = "A very fast and expressive template engine." -optional = false +optional = true python-versions = ">=3.7" files = [ {file = "jinja2-3.1.4-py3-none-any.whl", hash = "sha256:bc5dd2abb727a5319567b7a813e6a2e7318c39f4f487cfe6c89c6f9c7d25197d"}, @@ -814,7 +814,7 @@ testing = ["coverage", "pytest", "pytest-cov", "pytest-regressions"] name = "markupsafe" version = "2.1.5" description = "Safely add untrusted strings to HTML/XML markup." -optional = false +optional = true python-versions = ">=3.7" files = [ {file = "MarkupSafe-2.1.5-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:a17a92de5231666cfbe003f0e4b9b3a7ae3afb1ec2845aadc2bacc93ff85febc"}, @@ -1557,7 +1557,6 @@ files = [ {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:69b023b2b4daa7548bcfbd4aa3da05b3a74b772db9e23b982788168117739938"}, {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:81e0b275a9ecc9c0c0c07b4b90ba548307583c125f54d5b6946cfee6360c733d"}, {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba336e390cd8e4d1739f42dfe9bb83a3cc2e80f567d8805e11b46f4a943f5515"}, - {file = "PyYAML-6.0.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:326c013efe8048858a6d312ddd31d56e468118ad4cdeda36c719bf5bb6192290"}, {file = "PyYAML-6.0.1-cp310-cp310-win32.whl", hash = "sha256:bd4af7373a854424dabd882decdc5579653d7868b8fb26dc7d0e99f823aa5924"}, {file = "PyYAML-6.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:fd1592b3fdf65fff2ad0004b5e363300ef59ced41c2e6b3a99d4089fa8c5435d"}, {file = "PyYAML-6.0.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:6965a7bc3cf88e5a1c3bd2e0b5c22f8d677dc88a455344035f03399034eb3007"}, @@ -1565,16 +1564,8 @@ files = [ {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:42f8152b8dbc4fe7d96729ec2b99c7097d656dc1213a3229ca5383f973a5ed6d"}, {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:062582fca9fabdd2c8b54a3ef1c978d786e0f6b3a1510e0ac93ef59e0ddae2bc"}, {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d2b04aac4d386b172d5b9692e2d2da8de7bfb6c387fa4f801fbf6fb2e6ba4673"}, - {file = "PyYAML-6.0.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e7d73685e87afe9f3b36c799222440d6cf362062f78be1013661b00c5c6f678b"}, {file = "PyYAML-6.0.1-cp311-cp311-win32.whl", hash = "sha256:1635fd110e8d85d55237ab316b5b011de701ea0f29d07611174a1b42f1444741"}, {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"}, - {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"}, - {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"}, - {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"}, - {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"}, - {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"}, - {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"}, - {file = "PyYAML-6.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:0d3304d8c0adc42be59c5f8a4d9e3d7379e6955ad754aa9d6ab7a398b59dd1df"}, {file = "PyYAML-6.0.1-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:50550eb667afee136e9a77d6dc71ae76a44df8b3e51e41b77f6de2932bfe0f47"}, {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1fe35611261b29bd1de0070f0b2f47cb6ff71fa6595c077e42bd0c419fa27b98"}, {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:704219a11b772aea0d8ecd7058d0082713c3562b4e271b849ad7dc4a5c90c13c"}, @@ -1591,7 +1582,6 @@ files = [ {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a0cd17c15d3bb3fa06978b4e8958dcdc6e0174ccea823003a106c7d4d7899ac5"}, {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:28c119d996beec18c05208a8bd78cbe4007878c6dd15091efb73a30e90539696"}, {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7e07cbde391ba96ab58e532ff4803f79c4129397514e1413a7dc761ccd755735"}, - {file = "PyYAML-6.0.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:49a183be227561de579b4a36efbb21b3eab9651dd81b1858589f796549873dd6"}, {file = "PyYAML-6.0.1-cp38-cp38-win32.whl", hash = "sha256:184c5108a2aca3c5b3d3bf9395d50893a7ab82a38004c8f61c258d4428e80206"}, {file = "PyYAML-6.0.1-cp38-cp38-win_amd64.whl", hash = "sha256:1e2722cc9fbb45d9b87631ac70924c11d3a401b2d7f410cc0e3bbf249f2dca62"}, {file = "PyYAML-6.0.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:9eb6caa9a297fc2c2fb8862bc5370d0303ddba53ba97e71f08023b6cd73d16a8"}, @@ -1599,7 +1589,6 @@ files = [ {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5773183b6446b2c99bb77e77595dd486303b4faab2b086e7b17bc6bef28865f6"}, {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b786eecbdf8499b9ca1d697215862083bd6d2a99965554781d0d8d1ad31e13a0"}, {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bc1bf2925a1ecd43da378f4db9e4f799775d6367bdb94671027b73b393a7c42c"}, - {file = "PyYAML-6.0.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:04ac92ad1925b2cff1db0cfebffb6ffc43457495c9b3c39d3fcae417d7125dc5"}, {file = "PyYAML-6.0.1-cp39-cp39-win32.whl", hash = "sha256:faca3bdcf85b2fc05d06ff3fbc1f83e1391b3e724afa3feba7d13eeab355484c"}, {file = "PyYAML-6.0.1-cp39-cp39-win_amd64.whl", hash = "sha256:510c9deebc5c0225e8c96813043e62b680ba2f9c50a08d3724c7f28a747d1486"}, {file = "PyYAML-6.0.1.tar.gz", hash = "sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43"}, @@ -1877,6 +1866,17 @@ files = [ [package.extras] test = ["flake8", "mypy", "pytest"] +[[package]] +name = "sphinxcontrib-mermaid" +version = "0.9.2" +description = "Mermaid diagrams in yours Sphinx powered docs" +optional = true +python-versions = ">=3.7" +files = [ + {file = "sphinxcontrib-mermaid-0.9.2.tar.gz", hash = "sha256:252ef13dd23164b28f16d8b0205cf184b9d8e2b714a302274d9f59eb708e77af"}, + {file = "sphinxcontrib_mermaid-0.9.2-py3-none-any.whl", hash = "sha256:6795a72037ca55e65663d2a2c1a043d636dc3d30d418e56dd6087d1459d98a5d"}, +] + [[package]] name = "sphinxcontrib-qthelp" version = "1.0.7" @@ -2044,10 +2044,10 @@ numpy = ["numpy (>=1.9)"] web = ["wslink (>=1.0.4)"] [extras] -docs = ["Sphinx", "docutils", "myst-parser", "numpydoc", "readthedocs-sphinx-search", "sphinx-autodoc-typehints", "sphinx-book-theme"] +docs = ["Sphinx", "docutils", "myst-parser", "numpydoc", "readthedocs-sphinx-search", "sphinx-autodoc-typehints", "sphinx-book-theme", "sphinxcontrib-mermaid"] examples = ["cma"] [metadata] lock-version = "2.0" python-versions = ">=3.10,<3.12" -content-hash = "07105447f85d22c61c0766e22959a4f3428f9cfa20e7499aae4017fa993fdca3" +content-hash = "c9747151a346aa7c37c83f66c82b8283aa4ed796788d545d65378b7e8a117260" diff --git a/pyproject.toml b/pyproject.toml index 9851cafc9..bbd4e9ce4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,7 @@ Sphinx = {version = "^6.1", optional = true} sphinx-book-theme = {version = "^1.0", optional = true} readthedocs-sphinx-search = {version = ">=0.1.1,<0.4.0", optional = true} sphinx-autodoc-typehints = {version = "^1.21", optional = true} +sphinxcontrib-mermaid = {version = "^0.9.2", optional = true} myst-parser = {version = "^1.0", optional = true} numpydoc = {version = "^1.3.1", optional = true} docutils = {version = "^0.18", optional = true} @@ -73,6 +74,7 @@ docs = [ "sphinx-book-theme", "readthedocs-sphinx-search", "sphinx-autodoc-typehints", + "sphinxcontrib-mermaid", "myst-parser", "numpydoc", "docutils", From a036bc1355656baf66c1d182102e91b64504daf2 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Tue, 25 Jun 2024 07:42:10 -0500 Subject: [PATCH 112/134] type: reset functions --- .../_reset_ghost_vector_or_scalar.py | 26 +++++++++++++------ 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/elastica/reset_functions_for_block_structure/_reset_ghost_vector_or_scalar.py b/elastica/reset_functions_for_block_structure/_reset_ghost_vector_or_scalar.py index 3bac67448..4af0cd59d 100644 --- a/elastica/reset_functions_for_block_structure/_reset_ghost_vector_or_scalar.py +++ b/elastica/reset_functions_for_block_structure/_reset_ghost_vector_or_scalar.py @@ -1,18 +1,24 @@ __doc__ = """Reset the ghost vectors or scalar variables using functions implemented in Numba""" +import numpy +from numpy.typing import NDArray from numba import njit @njit(cache=True) # type: ignore -def _reset_vector_ghost(input, ghost_idx, reset_value=0.0): +def _reset_vector_ghost( + input: NDArray[np.floating], + ghost_idx: NDArray[np.integer], + reset_value: np.floating = np.float64(0.0), +): """ This function resets the ghost of an input vector collection. Default reset value is 0.0. Parameters ---------- - input - ghost_idx - reset_value + input : NDArray[np.floating] + ghost_idx : NDArray[np.integer] + reset_value : np.floating Returns ------- @@ -40,15 +46,19 @@ def _reset_vector_ghost(input, ghost_idx, reset_value=0.0): @njit(cache=True) # type: ignore -def _reset_scalar_ghost(input, ghost_idx, reset_value=0.0): +def _reset_scalar_ghost( + input: NDArray[np.floating], + ghost_idx: NDArray[np.integer], + reset_value: np.floating = np.float64(0.0), +): """ This function resets the ghost of a scalar collection. Default reset value is 0.0. Parameters ---------- - input - ghost_idx - reset_value + input : NDArray[np.floating] + ghost_idx : NDArray[np.integer] + reset_value : np.floating Returns ------- From 1e395ab59d46b09a39ae09cef5d4af6badabec0c Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Tue, 25 Jun 2024 07:45:27 -0500 Subject: [PATCH 113/134] revert: catenary case --- examples/CatenaryCase/catenary.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/examples/CatenaryCase/catenary.py b/examples/CatenaryCase/catenary.py index 6a2d9e88e..d6fd787b0 100644 --- a/examples/CatenaryCase/catenary.py +++ b/examples/CatenaryCase/catenary.py @@ -1,6 +1,5 @@ import numpy as np from elastica import * -from elastica.typing import SystemType, SystemCollectionType, SymplecticStepperProtocol from post_processing import ( plot_video, @@ -12,7 +11,7 @@ class CatenarySimulator(BaseSystemCollection, Constraints, Forcing, Damping, Cal pass -catenary_sim: SystemCollectionType = CatenarySimulator() +catenary_sim = CatenarySimulator() final_time = 10 damping_constant = 0.3 time_step = 1e-4 @@ -38,7 +37,7 @@ class CatenarySimulator(BaseSystemCollection, Constraints, Forcing, Damping, Cal poisson_ratio = 0.5 shear_modulus = E / (poisson_ratio + 1.0) -base_rod: SystemType = CosseratRod.straight_rod( +base_rod = CosseratRod.straight_rod( n_elem, start, direction, @@ -77,11 +76,11 @@ class CatenaryCallBack(CallBackBaseClass): """ def __init__(self, step_skip: int, callback_params: dict): - super().__init__() + CallBackBaseClass.__init__(self) self.every = step_skip self.callback_params = callback_params - def make_callback(self, system: SystemType, time: float, current_step: int) -> None: + def make_callback(self, system, time, current_step: int): if current_step % self.every == 0: @@ -94,7 +93,7 @@ def make_callback(self, system: SystemType, time: float, current_step: int) -> N return -pp_list: dict = defaultdict(list) +pp_list = defaultdict(list) catenary_sim.collect_diagnostics(base_rod).using( CatenaryCallBack, step_skip=step_skip, callback_params=pp_list ) @@ -103,7 +102,7 @@ def make_callback(self, system: SystemType, time: float, current_step: int) -> N catenary_sim.finalize() -timestepper: SymplecticStepperProtocol = PositionVerlet() +timestepper = PositionVerlet() integrate(timestepper, catenary_sim, final_time, total_steps) position = np.array(pp_list["position"]) From 256d85ac9037649dd4a31c41f7d8f1a59444a79f Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Tue, 25 Jun 2024 07:57:45 -0500 Subject: [PATCH 114/134] type: system factory function np.floating -> float --- elastica/rigidbody/cylinder.py | 24 +++---- elastica/rigidbody/sphere.py | 14 ++--- elastica/rod/cosserat_rod.py | 112 ++++++++++++++++----------------- elastica/surface/plane.py | 5 +- 4 files changed, 79 insertions(+), 76 deletions(-) diff --git a/elastica/rigidbody/cylinder.py b/elastica/rigidbody/cylinder.py index b9ff1e0a3..7c24e6da5 100644 --- a/elastica/rigidbody/cylinder.py +++ b/elastica/rigidbody/cylinder.py @@ -16,21 +16,21 @@ def __init__( start: NDArray[np.floating], direction: NDArray[np.floating], normal: NDArray[np.floating], - base_length: np.floating, - base_radius: np.floating, - density: np.floating, + base_length: float, + base_radius: float, + density: float, ) -> None: """ Rigid body cylinder initializer. Parameters ---------- - start - direction - normal - base_length - base_radius - density + start : NDArray[np.floating] + direction : NDArray[np.floating] + normal : NDArray[np.floating] + base_length : float + base_radius : float + density : float """ def _check_array_size( @@ -66,9 +66,9 @@ def _check_lower_bound( normal = normal.reshape((3, 1)) tangents = direction.reshape((3, 1)) binormal = _batch_cross(tangents, normal) - self.radius = base_radius - self.length = base_length - self.density = density + self.radius = np.float64(base_radius) + self.length = np.float64(base_length) + self.density = np.float64(density) dim: int = MaxDimension.value() diff --git a/elastica/rigidbody/sphere.py b/elastica/rigidbody/sphere.py index 7c0e7d803..38aaf23d9 100644 --- a/elastica/rigidbody/sphere.py +++ b/elastica/rigidbody/sphere.py @@ -14,17 +14,17 @@ class Sphere(RigidBodyBase): def __init__( self, center: NDArray[np.floating], - base_radius: np.floating, - density: np.floating, + base_radius: float, + density: float, ) -> None: """ Rigid body sphere initializer. Parameters ---------- - center - base_radius - density + center : NDArray[np.floating] + base_radius : float + density : float """ super().__init__() @@ -41,8 +41,8 @@ def __init__( assert base_radius > 0.0, "base_radius must be positive" assert density > 0.0, "density must be positive" - self.radius = base_radius - self.density = density + self.radius = np.float64(base_radiu) + self.density = np.float64(density) self.length = np.float64(2 * base_radius) # This is for a rigid body cylinder self.volume = np.float64(4.0 / 3.0 * np.pi * base_radius**3) diff --git a/elastica/rod/cosserat_rod.py b/elastica/rod/cosserat_rod.py index 868e9a912..aa36ca5da 100644 --- a/elastica/rod/cosserat_rod.py +++ b/elastica/rod/cosserat_rod.py @@ -79,73 +79,73 @@ class CosseratRod(RodBase, KnotTheory): ---------- n_elems: int The number of elements of the rod. - position_collection: numpy.ndarray + position_collection: NDArray[np.floating] 2D (dim, n_nodes) array containing data with 'float' type. Array containing node position vectors. - velocity_collection: numpy.ndarray + velocity_collection: NDArray[np.floating] 2D (dim, n_nodes) array containing data with 'float' type. Array containing node velocity vectors. - acceleration_collection: numpy.ndarray + acceleration_collection: NDArray[np.floating] 2D (dim, n_nodes) array containing data with 'float' type. Array containing node acceleration vectors. - omega_collection: numpy.ndarray + omega_collection: NDArray[np.floating] 2D (dim, n_elems) array containing data with 'float' type. Array containing element angular velocity vectors. - alpha_collection: numpy.ndarray + alpha_collection: NDArray[np.floating] 2D (dim, n_elems) array containing data with 'float' type. Array contining element angular acceleration vectors. - director_collection: numpy.ndarray + director_collection: NDArray[np.floating] 3D (dim, dim, n_elems) array containing data with 'float' type. Array containing element director matrices. - rest_lengths: numpy.ndarray + rest_lengths: NDArray[np.floating] 1D (n_elems) array containing data with 'float' type. Rod element lengths at rest configuration. - density: numpy.ndarray + density: NDArray[np.floating] 1D (n_elems) array containing data with 'float' type. Rod elements densities. - volume: numpy.ndarray + volume: NDArray[np.floating] 1D (n_elems) array containing data with 'float' type. Rod element volumes. - mass: numpy.ndarray + mass: NDArray[np.floating] 1D (n_nodes) array containing data with 'float' type. Rod node masses. Note that masses are stored on the nodes, not on elements. - mass_second_moment_of_inertia: numpy.ndarray + mass_second_moment_of_inertia: NDArray[np.floating] 3D (dim, dim, n_elems) array containing data with 'float' type. Rod element mass second moment of interia. - inv_mass_second_moment_of_inertia: numpy.ndarray + inv_mass_second_moment_of_inertia: NDArray[np.floating] 3D (dim, dim, n_elems) array containing data with 'float' type. Rod element inverse mass moment of inertia. - rest_voronoi_lengths: numpy.ndarray + rest_voronoi_lengths: NDArray[np.floating] 1D (n_voronoi) array containing data with 'float' type. Rod lengths on the voronoi domain at the rest configuration. - internal_forces: numpy.ndarray + internal_forces: NDArray[np.floating] 2D (dim, n_nodes) array containing data with 'float' type. Rod node internal forces. Note that internal forces are stored on the node, not on elements. - internal_torques: numpy.ndarray + internal_torques: NDArray[np.floating] 2D (dim, n_elems) array containing data with 'float' type. Rod element internal torques. - external_forces: numpy.ndarray + external_forces: NDArray[np.floating] 2D (dim, n_nodes) array containing data with 'float' type. External forces acting on rod nodes. - external_torques: numpy.ndarray + external_torques: NDArray[np.floating] 2D (dim, n_elems) array containing data with 'float' type. External torques acting on rod elements. - lengths: numpy.ndarray + lengths: NDArray[np.floating] 1D (n_elems) array containing data with 'float' type. Rod element lengths. - tangents: numpy.ndarray + tangents: NDArray[np.floating] 2D (dim, n_elems) array containing data with 'float' type. Rod element tangent vectors. - radius: numpy.ndarray + radius: NDArray[np.floating] 1D (n_elems) array containing data with 'float' type. Rod element radius. - dilatation: numpy.ndarray + dilatation: NDArray[np.floating] 1D (n_elems) array containing data with 'float' type. Rod element dilatation. - voronoi_dilatation: numpy.ndarray + voronoi_dilatation: NDArray[np.floating] 1D (n_voronoi) array containing data with 'float' type. Rod dilatation on voronoi domain. - dilatation_rate: numpy.ndarray + dilatation_rate: NDArray[np.floating] 1D (n_elems) array containing data with 'float' type. Rod element dilatation rates. """ @@ -253,12 +253,12 @@ def straight_rod( start: NDArray[np.floating], direction: NDArray[np.floating], normal: NDArray[np.floating], - base_length: np.floating, - base_radius: np.floating, - density: np.floating, + base_length: float, + base_radius: float, + density: float, *, nu: Optional[np.floating] = None, - youngs_modulus: np.floating, + youngs_modulus: float, **kwargs: Any, ) -> Self: """ @@ -276,21 +276,21 @@ def straight_rod( n_elements : int Number of element. Must be greater than 3. Generally recommended to start with 40-50, and adjust the resolution. - start : NDArray[3, float] + start : NDArray[np.floating] Starting coordinate in 3D - direction : NDArray[3, float] + direction : NDArray[np.floating] Direction of the rod in 3D - normal : NDArray[3, float] + normal : NDArray[np.floating] Normal vector of the rod in 3D - base_length : np.floating + base_length : float Total length of the rod - base_radius : np.floating + base_radius : float Uniform radius of the rod - density : np.floating + density : float Density of the rod - nu : np.floating + nu : float Damping coefficient for Rayleigh damping - youngs_modulus : np.floating + youngs_modulus : float Young's modulus **kwargs : dict, optional The "position" and/or "directors" can be overrided by passing "position" and "directors" argument. Remember, the shape of the "position" is (3,n_elements+1) and the shape of the "directors" is (3,3,n_elements). @@ -349,10 +349,10 @@ def straight_rod( n_elements, direction, normal, - base_length, - base_radius, - density, - youngs_modulus, + np.float64(base_length), + np.float64(base_radius), + np.float64(density), + np.float64(youngs_modulus), rod_origin_position=start, ring_rod_flag=ring_rod_flag, **kwargs, @@ -401,12 +401,12 @@ def ring_rod( ring_center_position: NDArray[np.floating], direction: NDArray[np.floating], normal: NDArray[np.floating], - base_length: np.floating, - base_radius: np.floating, - density: np.floating, + base_length: float, + base_radius: float, + density: float, *, - nu: Optional[np.floating] = None, - youngs_modulus: np.floating, + nu: Optional[float] = None, + youngs_modulus: float, **kwargs: Any, ) -> Self: """ @@ -423,21 +423,21 @@ def ring_rod( ---------- n_elements : int Number of element. Must be greater than 3. Generarally recommended to start with 40-50, and adjust the resolution. - ring_center_position : NDArray[3, float] + ring_center_position : NDArray[np.floating] Center coordinate for ring rod in 3D - direction : NDArray[3, float] + direction : NDArray[np.floating] Direction of the rod in 3D - normal : NDArray[3, float] + normal : NDArray[np.floating] Normal vector of the rod in 3D - base_length : np.floating + base_length : float Total length of the rod - base_radius : np.floating + base_radius : float Uniform radius of the rod - density : np.floating + density : float Density of the rod - nu : np.floating + nu : float | None Damping coefficient for Rayleigh damping - youngs_modulus : np.floating + youngs_modulus : float Young's modulus **kwargs : dict, optional The "position" and/or "directors" can be overrided by passing "position" and "directors" argument. Remember, the shape of the "position" is (3,n_elements+1) and the shape of the "directors" is (3,3,n_elements). @@ -497,10 +497,10 @@ def ring_rod( n_elements, direction, normal, - base_length, - base_radius, - density, - youngs_modulus, + np.float64(base_length), + np.float64(base_radius), + np.float64(density), + np.float64(youngs_modulus), rod_origin_position=ring_center_position, ring_rod_flag=ring_rod_flag, **kwargs, diff --git a/elastica/surface/plane.py b/elastica/surface/plane.py index 7f17537f0..99b8525d8 100644 --- a/elastica/surface/plane.py +++ b/elastica/surface/plane.py @@ -2,11 +2,14 @@ from elastica.surface.surface_base import SurfaceBase import numpy as np +from numpy.typing import NDArray from elastica.utils import Tolerance class Plane(SurfaceBase): - def __init__(self, plane_origin: np.ndarray, plane_normal: np.ndarray): + def __init__( + self, plane_origin: NDArray[np.floating], plane_normal: NDArray[np.floating] + ): """ Plane surface initializer. From 210d172b0e23991b3acfea014afb5b46d70ff368 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Tue, 25 Jun 2024 07:58:37 -0500 Subject: [PATCH 115/134] type: feature API np.floating -> float --- elastica/_calculus.py | 4 +-- elastica/boundary_conditions.py | 27 +++++---------- elastica/contact_forces.py | 60 ++++++++++++++++----------------- elastica/dissipation.py | 2 +- elastica/external_forces.py | 52 +++++++++++++++------------- elastica/interaction.py | 46 ++++++++++++++----------- elastica/joint.py | 32 +++++++++--------- 7 files changed, 112 insertions(+), 111 deletions(-) diff --git a/elastica/_calculus.py b/elastica/_calculus.py index edd7afdee..1639eafd2 100644 --- a/elastica/_calculus.py +++ b/elastica/_calculus.py @@ -4,7 +4,7 @@ from numpy import zeros, empty from numpy.typing import NDArray import numba -from numba import njit, float64 +from numba import njit from elastica.reset_functions_for_block_structure._reset_ghost_vector_or_scalar import ( _reset_vector_ghost, ) @@ -22,7 +22,7 @@ def _get_zero_array(dim: int, ndim: int) -> Union[float, NDArray[np.floating], N @njit(cache=True) # type: ignore -def _trapezoidal(array_collection: NDArray[np.float64]) -> NDArray[np.float64]: +def _trapezoidal(array_collection: NDArray[np.floating]) -> NDArray[np.floating]: """ Simple trapezoidal quadrature rule with zero at end-points, in a dimension agnostic way diff --git a/elastica/boundary_conditions.py b/elastica/boundary_conditions.py index 6b25b8653..66680e2d7 100644 --- a/elastica/boundary_conditions.py +++ b/elastica/boundary_conditions.py @@ -123,15 +123,6 @@ def constrain_rates( pass -class FreeRod(FreeBC): - # Please clear this part beyond version 0.3.0 - """Deprecated 0.2.1: Same implementation as FreeBC""" - warnings.warn( - "FreeRod is deprecated and renamed to FreeBC. The deprecated name will be removed in the future.", - DeprecationWarning, - ) - - class OneEndFixedBC(ConstraintBase): """ This boundary condition class fixes one end of the rod. Currently, @@ -702,9 +693,9 @@ def __init__( position_end: NDArray[np.floating], director_start: NDArray[np.floating], director_end: NDArray[np.floating], - twisting_time: np.floating, - slack: np.floating, - number_of_rotations: np.floating, + twisting_time: float, + slack: float, + number_of_rotations: float, **kwargs: Any, ) -> None: """ @@ -734,12 +725,12 @@ def __init__( Number of rotations applied to rod. """ super().__init__(**kwargs) - self.twisting_time = twisting_time + self.twisting_time = np.float64(twisting_time) - angel_vel_scalar = ( - 2.0 * number_of_rotations * np.pi / self.twisting_time - ) / 2.0 - shrink_vel_scalar = slack / (self.twisting_time * 2.0) + angel_vel_scalar = np.float64( + (2.0 * number_of_rotations * np.pi / self.twisting_time) / 2.0 + ) + shrink_vel_scalar = np.float64(slack / (self.twisting_time * 2.0)) direction = (position_end - position_start) / np.linalg.norm( position_end - position_start @@ -751,7 +742,7 @@ def __init__( self.ang_vel = angel_vel_scalar * direction self.shrink_vel = shrink_vel_scalar * direction - theta = number_of_rotations * np.pi + theta = np.float64(number_of_rotations * np.pi) self.final_start_directors = ( _get_rotation_matrix(theta, direction.reshape(3, 1)).reshape(3, 3) diff --git a/elastica/contact_forces.py b/elastica/contact_forces.py index 22adcb3e9..f2ee9ca78 100644 --- a/elastica/contact_forces.py +++ b/elastica/contact_forces.py @@ -196,8 +196,8 @@ class RodCylinderContact(NoContact): def __init__( self, - k: np.floating, - nu: np.floating, + k: float, + nu: float, velocity_damping_coefficient: float = 0.0, friction_coefficient: float = 0.0, ) -> None: @@ -216,10 +216,10 @@ def __init__( For Coulombic friction coefficient for rigid-body and rod contact. """ super(RodCylinderContact, self).__init__() - self.k = k - self.nu = nu - self.velocity_damping_coefficient = velocity_damping_coefficient - self.friction_coefficient = friction_coefficient + self.k = np.float64(k) + self.nu = np.float64(nu) + self.velocity_damping_coefficient = np.float64(velocity_damping_coefficient) + self.friction_coefficient = np.float64(friction_coefficient) @property def _allowed_system_two(self) -> list[Type]: @@ -287,7 +287,7 @@ class RodSelfContact(NoContact): """ - def __init__(self, k: np.floating, nu: np.floating) -> None: + def __init__(self, k: float, nu: float) -> None: """ Parameters @@ -298,8 +298,8 @@ def __init__(self, k: np.floating, nu: np.floating) -> None: Contact damping constant. """ super(RodSelfContact, self).__init__() - self.k = k - self.nu = nu + self.k = np.float64(k) + self.nu = np.float64(nu) def _check_systems_validity( self, @@ -364,8 +364,8 @@ class RodSphereContact(NoContact): def __init__( self, - k: np.floating, - nu: np.floating, + k: float, + nu: float, velocity_damping_coefficient: float = 0.0, friction_coefficient: float = 0.0, ) -> None: @@ -383,9 +383,9 @@ def __init__( For Coulombic friction coefficient for rigid-body and rod contact. """ super(RodSphereContact, self).__init__() - self.k = k - self.nu = nu - self.velocity_damping_coefficient = velocity_damping_coefficient + self.k = np.float64(k) + self.nu = np.float64(nu) + self.velocity_damping_coefficient = np.float64(velocity_damping_coefficient) self.friction_coefficient = np.float64(friction_coefficient) @property @@ -465,8 +465,8 @@ class RodPlaneContact(NoContact): def __init__( self, - k: np.floating, - nu: np.floating, + k: float, + nu: float, ) -> None: """ Parameters @@ -477,9 +477,9 @@ def __init__( Contact damping constant. """ super(RodPlaneContact, self).__init__() - self.k = k - self.nu = nu - self.surface_tol = 1e-4 + self.k = np.float64(k) + self.nu = np.float64(nu) + self.surface_tol = np.float64(1.0e-4) @property def _allowed_system_two(self) -> list[Type]: @@ -535,9 +535,9 @@ class RodPlaneContactWithAnisotropicFriction(NoContact): def __init__( self, - k: np.floating, - nu: np.floating, - slip_velocity_tol: np.floating, + k: float, + nu: float, + slip_velocity_tol: float, static_mu_array: NDArray[np.floating], kinetic_mu_array: NDArray[np.floating], ) -> None: @@ -558,9 +558,9 @@ def __init__( [forward, backward, sideways] kinetic friction coefficients. """ super(RodPlaneContactWithAnisotropicFriction, self).__init__() - self.k = k - self.nu = nu - self.surface_tol = 1e-4 + self.k = np.float64(k) + self.nu = np.float64(nu) + self.surface_tol = np.float64(1.0e-4) self.slip_velocity_tol = slip_velocity_tol ( self.static_mu_forward, @@ -635,8 +635,8 @@ class CylinderPlaneContact(NoContact): def __init__( self, - k: np.floating, - nu: np.floating, + k: float, + nu: float, ) -> None: """ Parameters @@ -647,9 +647,9 @@ def __init__( Contact damping constant. """ super(CylinderPlaneContact, self).__init__() - self.k = k - self.nu = nu - self.surface_tol = 1e-4 + self.k = np.float64(k) + self.nu = np.float64(nu) + self.surface_tol = np.float64(1.0e-4) @property def _allowed_system_one(self) -> list[Type]: diff --git a/elastica/dissipation.py b/elastica/dissipation.py index 2cc532918..7a5718aa4 100644 --- a/elastica/dissipation.py +++ b/elastica/dissipation.py @@ -120,7 +120,7 @@ class AnalyticalLinearDamper(DamperBase): """ def __init__( - self, damping_constant: np.floating, time_step: np.floating, **kwargs: Any + self, damping_constant: float, time_step: float, **kwargs: Any ) -> None: """ Analytical linear damper initializer diff --git a/elastica/external_forces.py b/elastica/external_forces.py index ae85a3569..e9b5b7238 100644 --- a/elastica/external_forces.py +++ b/elastica/external_forces.py @@ -34,7 +34,7 @@ def __init__(self) -> None: """ pass - def apply_forces(self, system: S, time: float = 0.0) -> None: + def apply_forces(self, system: S, time: np.floating = np.float64(0.0)) -> None: """Apply forces to a rod-like object. In NoForces class, this routine simply passes. @@ -49,7 +49,7 @@ def apply_forces(self, system: S, time: float = 0.0) -> None: """ pass - def apply_torques(self, system: S, time: float = 0.0) -> None: + def apply_torques(self, system: S, time: np.floating = np.float64(0.0)) -> None: """Apply torques to a rod-like object. In NoForces class, this routine simply passes. @@ -91,7 +91,7 @@ def __init__( self.acc_gravity = acc_gravity def apply_forces( - self, system: "RodType | RigidBodyType", time: float = 0.0 + self, system: "RodType | RigidBodyType", time: np.floating = np.float64(0.0) ) -> None: self.compute_gravity_forces( self.acc_gravity, system.mass, system.external_forces @@ -160,10 +160,10 @@ def __init__( self.start_force = start_force self.end_force = end_force assert ramp_up_time > 0.0 - self.ramp_up_time = ramp_up_time + self.ramp_up_time = np.float64(ramp_up_time) def apply_forces( - self, system: "RodType | RigidBodyType", time: float = 0.0 + self, system: "RodType | RigidBodyType", time: np.floating = np.float64(0.0) ) -> None: self.compute_end_point_forces( system.external_forces, @@ -179,7 +179,7 @@ def compute_end_point_forces( external_forces: NDArray[np.floating], start_force: NDArray[np.floating], end_force: NDArray[np.floating], - time: float, + time: np.floating, ramp_up_time: np.floating, ) -> None: """ @@ -234,7 +234,7 @@ def __init__( self.torque = torque * direction def apply_torques( - self, system: "RodType | RigidBodyType", time: float = 0.0 + self, system: "RodType | RigidBodyType", time: np.floating = np.float64(0.0) ) -> None: n_elems = system.n_elems torque_on_one_element = ( @@ -273,7 +273,9 @@ def __init__( super(UniformForces, self).__init__() self.force = (force * direction).reshape(3, 1) - def apply_forces(self, rod: "RodType | RigidBodyType", time: float = 0.0) -> None: + def apply_forces( + self, rod: "RodType | RigidBodyType", time: np.floating = np.float64(0.0) + ) -> None: force_on_one_element = self.force / rod.n_elems rod.external_forces += force_on_one_element @@ -310,14 +312,14 @@ class MuscleTorques(NoForces): def __init__( self, - base_length: np.floating, + base_length: float, # TODO: Is this necessary? b_coeff: NDArray[np.floating], - period: np.floating, - wave_number: np.floating, - phase_shift: np.floating, + period: float, + wave_number: float, + phase_shift: float, direction: NDArray[np.floating], rest_lengths: NDArray[np.floating], - ramp_up_time: np.floating, + ramp_up_time: float, with_spline: bool = False, ) -> None: """ @@ -346,12 +348,12 @@ def __init__( super(MuscleTorques, self).__init__() self.direction = direction # Direction torque applied - self.angular_frequency = 2.0 * np.pi / period - self.wave_number = wave_number - self.phase_shift = phase_shift + self.angular_frequency = np.float64(2.0 * np.pi / period) + self.wave_number = np.float64(wave_number) + self.phase_shift = np.float64(phase_shift) assert ramp_up_time > 0.0 - self.ramp_up_time = ramp_up_time + self.ramp_up_time = np.float64(ramp_up_time) # s is the position of nodes on the rod, we go from node=1 to node=nelem-1, because there is no # torques applied by first and last node on elements. Reason is that we cannot apply torque in an @@ -370,7 +372,9 @@ def __init__( else: self.my_spline = np.full_like(self.s, fill_value=1.0) - def apply_torques(self, rod: "RodType | RigidBodyType", time: float = 0.0) -> None: + def apply_torques( + self, rod: "RodType | RigidBodyType", time: np.floating = np.float64(0.0) + ) -> None: self.compute_muscle_torques( time, self.my_spline, @@ -501,8 +505,8 @@ class EndpointForcesSinusoidal(NoForces): def __init__( self, - start_force_mag: np.floating, - end_force_mag: np.floating, + start_force_mag: float, + end_force_mag: float, ramp_up_time: float = 0.0, tangent_direction: NDArray[np.floating] = np.array([0, 0, 1]), normal_direction: NDArray[np.floating] = np.array([0, 1, 0]), @@ -526,18 +530,18 @@ def __init__( """ super(EndpointForcesSinusoidal, self).__init__() # Start force - self.start_force_mag = start_force_mag - self.end_force_mag = end_force_mag + self.start_force_mag = np.float64(start_force_mag) + self.end_force_mag = np.float64(end_force_mag) # Applied force directions self.normal_direction = normal_direction self.roll_direction = np.cross(normal_direction, tangent_direction) assert ramp_up_time >= 0.0 - self.ramp_up_time = ramp_up_time + self.ramp_up_time = np.float64(ramp_up_time) def apply_forces( - self, system: "RodType | RigidBodyType", time: float = 0.0 + self, system: "RodType | RigidBodyType", time: np.floating = np.float64(0.0) ) -> None: if time < self.ramp_up_time: diff --git a/elastica/interaction.py b/elastica/interaction.py index 24ecb92e7..70ffcd34b 100644 --- a/elastica/interaction.py +++ b/elastica/interaction.py @@ -46,8 +46,8 @@ class InteractionPlane(NoForces): def __init__( self, - k: np.floating, - nu: np.floating, + k: float, + nu: float, plane_origin: NDArray[np.floating], plane_normal: NDArray[np.floating], ) -> None: @@ -66,13 +66,15 @@ def __init__( 2D (dim, 1) array containing data with 'float' type. The normal vector of the plane. """ - self.k = k - self.nu = nu + self.k = np.float64(k) + self.nu = np.float64(nu) + self.surface_tol = np.float64(1e-4) self.plane_origin = plane_origin.reshape(3, 1) self.plane_normal = plane_normal.reshape(3) - self.surface_tol = 1e-4 - def apply_forces(self, system: RodType, time: float = 0.0) -> None: + def apply_forces( + self, system: RodType, time: np.floating = np.float64(0.0) + ) -> None: """ In the case of contact with the plane, this function computes the plane reaction force on the element. @@ -142,11 +144,11 @@ class AnisotropicFrictionalPlane(InteractionPlane): def __init__( self, - k: np.floating, - nu: np.floating, + k: float, + nu: float, plane_origin: NDArray[np.floating], plane_normal: NDArray[np.floating], - slip_velocity_tol: np.floating, + slip_velocity_tol: float, static_mu_array: NDArray[np.floating], kinetic_mu_array: NDArray[np.floating], ) -> None: @@ -174,7 +176,7 @@ def __init__( [forward, backward, sideways] kinetic friction coefficients. """ InteractionPlane.__init__(self, k, nu, plane_origin, plane_normal) - self.slip_velocity_tol = slip_velocity_tol + self.slip_velocity_tol = np.float64(slip_velocity_tol) ( self.static_mu_forward, self.static_mu_backward, @@ -189,7 +191,7 @@ def __init__( # kinetic and static friction should separate functions # for now putting them together to figure out common variables def apply_forces( - self, system: "RodType | RigidBodyType", time: float = 0.0 + self, system: "RodType | RigidBodyType", time: np.floating = np.float64(0.0) ) -> None: """ Call numba implementation to apply friction forces @@ -386,7 +388,7 @@ class SlenderBodyTheory(NoForces): """ - def __init__(self, dynamic_viscosity: np.floating) -> None: + def __init__(self, dynamic_viscosity: float) -> None: """ Parameters @@ -395,9 +397,11 @@ def __init__(self, dynamic_viscosity: np.floating) -> None: Dynamic viscosity of the fluid. """ super(SlenderBodyTheory, self).__init__() - self.dynamic_viscosity = dynamic_viscosity + self.dynamic_viscosity = np.float64(dynamic_viscosity) - def apply_forces(self, system: RodType, time: float = 0.0) -> None: + def apply_forces( + self, system: RodType, time: np.floating = np.float64(0.0) + ) -> None: """ This function applies hydrodynamic forces on body using the slender body theory given in @@ -425,18 +429,20 @@ def apply_forces(self, system: RodType, time: float = 0.0) -> None: class InteractionPlaneRigidBody(NoForces): def __init__( self, - k: np.floating, - nu: np.floating, + k: float, + nu: float, plane_origin: NDArray[np.floating], plane_normal: NDArray[np.floating], ) -> None: - self.k = k - self.nu = nu + self.k = np.float64(k) + self.nu = np.float64(nu) + self.surface_tol = np.float64(1e-4) self.plane_origin = plane_origin.reshape(3, 1) self.plane_normal = plane_normal.reshape(3) - self.surface_tol = 1e-4 - def apply_forces(self, system: RigidBodyType, time: float = 0.0) -> None: + def apply_forces( + self, system: RigidBodyType, time: np.floating = np.float64(0.0) + ) -> None: """ This function computes the plane force response on the rigid body, in the case of contact. Contact model given in Eqn 4.8 Gazzola et. al. RSoS 2018 paper diff --git a/elastica/joint.py b/elastica/joint.py index 812a754c9..038ac256d 100644 --- a/elastica/joint.py +++ b/elastica/joint.py @@ -1,7 +1,7 @@ __doc__ = """ Module containing joint classes to connect multiple rods together. """ __all__ = ["FreeJoint", "HingeJoint", "FixedJoint", "get_relative_rotation_two_systems"] -from typing import Optional +pass from elastica._rotations import _inv_rotate from elastica.typing import SystemType, RodType, ConnectionIndex, RigidBodyType @@ -32,7 +32,7 @@ class FreeJoint: # pass the k and nu for the forces # also the necessary rods for the joint # indices should be 0 or -1, we will provide wrappers for users later - def __init__(self, k: np.floating, nu: np.floating) -> None: + def __init__(self, k: float, nu: float) -> None: """ Parameters @@ -43,8 +43,8 @@ def __init__(self, k: np.floating, nu: np.floating) -> None: Damping coefficient of the joint. """ - self.k = k - self.nu = nu + self.k = np.float64(k) + self.nu = np.float64(nu) def apply_forces( self, @@ -142,9 +142,9 @@ class HingeJoint(FreeJoint): # TODO: IN WRAPPER COMPUTE THE NORMAL DIRECTION OR ASK USER TO GIVE INPUT, IF NOT THROW ERROR def __init__( self, - k: np.floating, - nu: np.floating, - kt: np.floating, + k: float, + nu: float, + kt: float, normal_direction: NDArray[np.floating], ) -> None: """ @@ -167,7 +167,7 @@ def __init__( self.normal_direction = normal_direction / np.linalg.norm(normal_direction) # additional in-plane constraint through restoring torque # stiffness of the restoring constraint -- tuned empirically - self.kt = kt + self.kt = np.float64(kt) # Apply force is same as free joint def apply_forces( @@ -236,11 +236,11 @@ class FixedJoint(FreeJoint): def __init__( self, - k: np.floating, - nu: np.floating, - kt: np.floating, - nut: np.floating = np.float64(0.0), - rest_rotation_matrix: Optional[NDArray[np.floating]] = None, + k: float, + nu: float, + kt: float, + nut: float = 0.0, + rest_rotation_matrix: NDArray[np.floating] | None = None, ) -> None: """ @@ -254,7 +254,7 @@ def __init__( Rotational stiffness coefficient of the joint. nut: float = 0. Rotational damping coefficient of the joint. - rest_rotation_matrix: np.array + rest_rotation_matrix: np.array | None 2D (3,3) array containing data with 'float' type. Rest 3x3 rotation matrix from system one to system two at the connected elements. If provided, the rest rotation matrix is enforced between the two systems throughout the simulation. @@ -265,8 +265,8 @@ def __init__( super().__init__(k, nu) # additional in-plane constraint through restoring torque # stiffness of the restoring constraint -- tuned empirically - self.kt = kt - self.nut = nut + self.kt = np.float64(kt) + self.nut = np.float64(nut) # TODO: compute the rest rotation matrix directly during initialization # as soon as systems (e.g. `rod_one` and `rod_two`) and indices (e.g. `index_one` and `index_two`) From e2c65f11526e4dff03f500c46d5bba2eb066fe0b Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Tue, 25 Jun 2024 08:04:40 -0500 Subject: [PATCH 116/134] fix CI: minor typo, deprecated functions, and mypy --- elastica/__init__.py | 2 - elastica/boundary_conditions.py | 10 ----- elastica/joint.py | 2 - .../_reset_ghost_vector_or_scalar.py | 2 +- elastica/rigidbody/cylinder.py | 38 ++++++++++++------- elastica/rigidbody/rigid_body.py | 4 ++ elastica/rigidbody/sphere.py | 17 ++++++--- elastica/rod/cosserat_rod.py | 8 ++-- 8 files changed, 46 insertions(+), 37 deletions(-) diff --git a/elastica/__init__.py b/elastica/__init__.py index f705615a4..21dfddb72 100644 --- a/elastica/__init__.py +++ b/elastica/__init__.py @@ -17,8 +17,6 @@ GeneralConstraint, FixedConstraint, HelicalBucklingBC, - FreeRod, - OneEndFixedRod, ) from elastica.external_forces import ( NoForces, diff --git a/elastica/boundary_conditions.py b/elastica/boundary_conditions.py index 66680e2d7..87ce61b28 100644 --- a/elastica/boundary_conditions.py +++ b/elastica/boundary_conditions.py @@ -1,6 +1,5 @@ __doc__ = """ Built-in boundary condition implementationss """ -import warnings from typing import Any, Optional, TypeVar, Generic import numpy as np @@ -238,15 +237,6 @@ def compute_constrain_rates( omega_collection[..., 0] = 0.0 -class OneEndFixedRod(OneEndFixedBC): - # Please clear this part beyond version 0.3.0 - """Deprecated 0.2.1: Same implementation as OneEndFixedBC""" - warnings.warn( - "OneEndFixedRod is deprecated and renamed to OneEndFixedBC. The deprecated name will be removed in the future.", - DeprecationWarning, - ) - - class GeneralConstraint(ConstraintBase): """ This boundary condition class allows the specified node/link to have a configurable constraint. diff --git a/elastica/joint.py b/elastica/joint.py index 038ac256d..2788b5019 100644 --- a/elastica/joint.py +++ b/elastica/joint.py @@ -1,8 +1,6 @@ __doc__ = """ Module containing joint classes to connect multiple rods together. """ __all__ = ["FreeJoint", "HingeJoint", "FixedJoint", "get_relative_rotation_two_systems"] -pass - from elastica._rotations import _inv_rotate from elastica.typing import SystemType, RodType, ConnectionIndex, RigidBodyType diff --git a/elastica/reset_functions_for_block_structure/_reset_ghost_vector_or_scalar.py b/elastica/reset_functions_for_block_structure/_reset_ghost_vector_or_scalar.py index 4af0cd59d..18e5d129c 100644 --- a/elastica/reset_functions_for_block_structure/_reset_ghost_vector_or_scalar.py +++ b/elastica/reset_functions_for_block_structure/_reset_ghost_vector_or_scalar.py @@ -1,6 +1,6 @@ __doc__ = """Reset the ghost vectors or scalar variables using functions implemented in Numba""" -import numpy +import numpy as np from numpy.typing import NDArray from numba import njit diff --git a/elastica/rigidbody/cylinder.py b/elastica/rigidbody/cylinder.py index 7c24e6da5..6dab8384c 100644 --- a/elastica/rigidbody/cylinder.py +++ b/elastica/rigidbody/cylinder.py @@ -1,6 +1,7 @@ __doc__ = """ Implementation of a rigid body cylinder. """ +from typing import TYPE_CHECKING import numpy as np from numpy.typing import NDArray @@ -33,7 +34,8 @@ def __init__( density : float """ - def _check_array_size( + # FIXME: Refactor + def assert_check_array_size( to_check: NDArray[np.floating], name: str, expected: int = 3 ) -> None: array_size = to_check.size @@ -42,27 +44,24 @@ def _check_array_size( f"Expected: {expected}, but got: {array_size}" ) - def _check_lower_bound( - to_check: np.floating, name: str, lower_bound: np.floating = np.float64(0.0) + # FIXME: Refactor + def assert_check_lower_bound( + to_check: float, name: str, lower_bound: float = 0.0 ) -> None: assert ( to_check > lower_bound ), f"Value for '{name}' ({to_check}) must be at lease {lower_bound}. " - _check_array_size(start, "start") - _check_array_size(direction, "direction") - _check_array_size(normal, "normal") + assert_check_array_size(start, "start") + assert_check_array_size(direction, "direction") + assert_check_array_size(normal, "normal") - _check_lower_bound(base_length, "base_length") - _check_lower_bound(base_radius, "base_radius") - _check_lower_bound(density, "density") + assert_check_lower_bound(base_length, "base_length") + assert_check_lower_bound(base_radius, "base_radius") + assert_check_lower_bound(density, "density") super().__init__() - # rigid body does not have elements it only has one node. We are setting n_elems to - # zero for only make code to work. _bootstrap_from_data requires n_elems to be define - self.n_elem: int = 1 - normal = normal.reshape((3, 1)) tangents = direction.reshape((3, 1)) binormal = _batch_cross(tangents, normal) @@ -113,3 +112,16 @@ def _check_lower_bound( self.director_collection[0, ...] = normal self.director_collection[1, ...] = binormal self.director_collection[2, ...] = tangents + + +if TYPE_CHECKING: + from .protocol import RigidBodyProtocol + + _: RigidBodyProtocol = Cylinder( + start=np.zeros(3), + direction=np.ones(3), + normal=np.ones(3), + base_length=1.0, + base_radius=1.0, + density=1.0, + ) diff --git a/elastica/rigidbody/rigid_body.py b/elastica/rigidbody/rigid_body.py index c6c478151..a7252e049 100644 --- a/elastica/rigidbody/rigid_body.py +++ b/elastica/rigidbody/rigid_body.py @@ -22,6 +22,10 @@ class RigidBodyBase(ABC): REQUISITE_MODULES: list[Type] = [] def __init__(self) -> None: + # rigid body does not have elements it only has one node. We are setting n_elems to + # make code to work. _bootstrap_from_data requires n_elems to be define + self.n_elems: int = 1 + self.n_nodes: int = 1 self.position_collection: NDArray[np.floating] self.velocity_collection: NDArray[np.floating] diff --git a/elastica/rigidbody/sphere.py b/elastica/rigidbody/sphere.py index 38aaf23d9..72c4481b8 100644 --- a/elastica/rigidbody/sphere.py +++ b/elastica/rigidbody/sphere.py @@ -1,6 +1,7 @@ __doc__ = """ Implementation of a sphere rigid body. """ +from typing import TYPE_CHECKING import numpy as np from numpy.typing import NDArray @@ -29,10 +30,6 @@ def __init__( super().__init__() - # rigid body does not have elements it only have one node. We are setting n_elems to - # zero for only make code to work. _bootstrap_from_data requires n_elems to be defined - self.n_elems: int = 1 - dim: int = MaxDimension.value() assert ( @@ -41,7 +38,7 @@ def __init__( assert base_radius > 0.0, "base_radius must be positive" assert density > 0.0, "density must be positive" - self.radius = np.float64(base_radiu) + self.radius = np.float64(base_radius) self.density = np.float64(density) self.length = np.float64(2 * base_radius) # This is for a rigid body cylinder @@ -82,3 +79,13 @@ def __init__( self.director_collection[0, ...] = normal self.director_collection[1, ...] = binormal self.director_collection[2, ...] = tangents + + +if TYPE_CHECKING: + from .protocol import RigidBodyProtocol + + _: RigidBodyProtocol = Sphere( + center=np.zeros(3), + base_radius=1.0, + density=1.0, + ) diff --git a/elastica/rod/cosserat_rod.py b/elastica/rod/cosserat_rod.py index aa36ca5da..ea79854a1 100644 --- a/elastica/rod/cosserat_rod.py +++ b/elastica/rod/cosserat_rod.py @@ -1135,8 +1135,8 @@ def _zeroed_out_external_forces_and_torques( np.zeros(3), np.array([0, 1, 0]), np.array([0, 0, 1]), - np.float64(1.0), - np.float64(0.1), - np.float64(1.0), - youngs_modulus=np.float64(1.0), + 1.0, + 0.1, + 1.0, + youngs_modulus=1.0, ) From d3ddc66c4e94e133dd7235606d01bb6a7c6c0a94 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Tue, 25 Jun 2024 09:29:08 -0500 Subject: [PATCH 117/134] type: remove ignore during timestepper, use type-cast --- elastica/timestepper/__init__.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/elastica/timestepper/__init__.py b/elastica/timestepper/__init__.py index 0fd32c751..d2d095365 100644 --- a/elastica/timestepper/__init__.py +++ b/elastica/timestepper/__init__.py @@ -1,6 +1,6 @@ __doc__ = """Timestepping utilities to be used with Rod and RigidBody classes""" -from typing import Tuple, List, Callable, Type, Any, overload +from typing import Tuple, List, Callable, Type, Any, overload, cast from elastica.typing import SystemType, SystemCollectionType, SteppersOperatorsType import numpy as np @@ -81,15 +81,17 @@ def integrate( assert final_time > 0.0, "Final time is negative!" assert n_steps > 0, "Number of integration steps is negative!" - dt = np.float_(float(final_time) / n_steps) - time = np.float_(restart_time) + dt = np.float64(float(final_time) / n_steps) + time = np.float64(restart_time) if is_system_a_collection(systems): + systems = cast(SystemCollectionType, systems) for i in tqdm(range(n_steps), disable=(not progress_bar)): - time = stepper.step(systems, time, dt) # type: ignore[arg-type] + time = stepper.step(systems, time, dt) else: + systems = cast(SystemType, systems) for i in tqdm(range(n_steps), disable=(not progress_bar)): - time = stepper.step_single_instance(systems, time, dt) # type: ignore[arg-type] + time = stepper.step_single_instance(systems, time, dt) print("Final time of simulation is : ", time) return float(time) From df8477a7db46973bc1b28a2ff93b7f6120cd4ccd Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Tue, 25 Jun 2024 09:30:40 -0500 Subject: [PATCH 118/134] bug: fix issue that boudnary condition was not properly beling applied --- elastica/boundary_conditions.py | 18 ++++++++++++------ elastica/modules/constraints.py | 8 +++++++- 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/elastica/boundary_conditions.py b/elastica/boundary_conditions.py index 87ce61b28..20c10177e 100644 --- a/elastica/boundary_conditions.py +++ b/elastica/boundary_conditions.py @@ -11,7 +11,7 @@ from elastica._linalg import _batch_matvec, _batch_matrix_transpose from elastica._rotations import _get_rotation_matrix -from elastica.typing import SystemType, RodType, RigidBodyType +from elastica.typing import SystemType, RodType, RigidBodyType, ConstrainingIndex S = TypeVar("S") @@ -34,18 +34,24 @@ class ConstraintBase(ABC, Generic[S]): """ _system: S - _constrained_position_idx: np.ndarray - _constrained_director_idx: np.ndarray + _constrained_position_idx: NDArray[np.integer] + _constrained_director_idx: NDArray[np.integer] - def __init__(self, *args: Any, **kwargs: Any) -> None: + def __init__( + self, + *args: Any, + constrained_position_idx: ConstrainingIndex, + constrained_director_idx: ConstrainingIndex, + **kwargs: Any, + ) -> None: """Initialize boundary condition""" try: self._system = kwargs["_system"] self._constrained_position_idx = np.array( - kwargs.get("constrained_position_idx", []), dtype=int + constrained_position_idx, dtype=np.int_ ) self._constrained_director_idx = np.array( - kwargs.get("constrained_director_idx", []), dtype=int + constrained_director_idx, dtype=np.int_ ) except KeyError: raise KeyError( diff --git a/elastica/modules/constraints.py b/elastica/modules/constraints.py index 344306b31..2ba907f0c 100644 --- a/elastica/modules/constraints.py +++ b/elastica/modules/constraints.py @@ -215,7 +215,13 @@ def instantiate(self, system: "RodType | RigidBodyType") -> ConstraintBase: ) try: bc = self._bc_cls( - *positions, *directors, *self._args, _system=system, **self._kwargs + *positions, + *directors, + *self._args, + _system=system, + constrained_position_idx=self.constrained_position_idx, + constrained_director_idx=self.constrained_director_idx, + **self._kwargs, ) return bc except (TypeError, IndexError): From 06cc1e0404496a52c42d64c37ce246508fbf1dfe Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Tue, 25 Jun 2024 09:45:59 -0500 Subject: [PATCH 119/134] refactor: rename rod->system for proper inheritance syntax --- elastica/boundary_conditions.py | 32 ++++++++++++++++---------------- elastica/dissipation.py | 16 ++++++++-------- elastica/external_forces.py | 16 ++++++++-------- tests/test_dissipation.py | 8 ++++---- 4 files changed, 36 insertions(+), 36 deletions(-) diff --git a/elastica/boundary_conditions.py b/elastica/boundary_conditions.py index 20c10177e..0a324f119 100644 --- a/elastica/boundary_conditions.py +++ b/elastica/boundary_conditions.py @@ -88,7 +88,7 @@ def constrain_values(self, system: S, time: np.floating) -> None: time : float The time of simulation. """ - pass + raise NotImplementedError @abstractmethod def constrain_rates(self, system: S, time: np.floating) -> None: @@ -104,7 +104,7 @@ def constrain_rates(self, system: S, time: np.floating) -> None: The time of simulation. """ - pass + raise NotImplementedError class FreeBC(ConstraintBase): @@ -750,28 +750,28 @@ def __init__( ) # rotation_matrix wants vectors 3,1 def constrain_values( - self, rod: "RodType | RigidBodyType", time: np.floating + self, system: "RodType | RigidBodyType", time: np.floating ) -> None: if time > self.twisting_time: - rod.position_collection[..., 0] = self.final_start_position - rod.position_collection[..., -1] = self.final_end_position + system.position_collection[..., 0] = self.final_start_position + system.position_collection[..., -1] = self.final_end_position - rod.director_collection[..., 0] = self.final_start_directors - rod.director_collection[..., -1] = self.final_end_directors + system.director_collection[..., 0] = self.final_start_directors + system.director_collection[..., -1] = self.final_end_directors def constrain_rates( - self, rod: "RodType | RigidBodyType", time: np.floating + self, system: "RodType | RigidBodyType", time: np.floating ) -> None: if time > self.twisting_time: - rod.velocity_collection[..., 0] = 0.0 - rod.omega_collection[..., 0] = 0.0 + system.velocity_collection[..., 0] = 0.0 + system.omega_collection[..., 0] = 0.0 - rod.velocity_collection[..., -1] = 0.0 - rod.omega_collection[..., -1] = 0.0 + system.velocity_collection[..., -1] = 0.0 + system.omega_collection[..., -1] = 0.0 else: - rod.velocity_collection[..., 0] = self.shrink_vel - rod.omega_collection[..., 0] = self.ang_vel + system.velocity_collection[..., 0] = self.shrink_vel + system.omega_collection[..., 0] = self.ang_vel - rod.velocity_collection[..., -1] = -self.shrink_vel - rod.omega_collection[..., -1] = -self.ang_vel + system.velocity_collection[..., -1] = -self.shrink_vel + system.omega_collection[..., -1] = -self.ang_vel diff --git a/elastica/dissipation.py b/elastica/dissipation.py index 7a5718aa4..f20c632e0 100644 --- a/elastica/dissipation.py +++ b/elastica/dissipation.py @@ -151,13 +151,13 @@ def __init__( * np.diagonal(self._system.inv_mass_second_moment_of_inertia).T ) - def dampen_rates(self, rod: RodType, time: np.floating) -> None: - rod.velocity_collection[:] = ( - rod.velocity_collection * self.translational_damping_coefficient + def dampen_rates(self, system: RodType, time: np.floating) -> None: + system.velocity_collection[:] = ( + system.velocity_collection * self.translational_damping_coefficient ) - rod.omega_collection[:] = rod.omega_collection * np.power( - self.rotational_damping_coefficient, rod.dilatation + system.omega_collection[:] = system.omega_collection * np.power( + self.rotational_damping_coefficient, system.dilatation ) @@ -240,12 +240,12 @@ def __init__(self, filter_order: int, **kwargs: Any) -> None: self.omega_filter_term = np.zeros_like(self._system.omega_collection) self.filter_function = _filter_function_periodic_condition - def dampen_rates(self, rod: RodType, time: np.floating) -> None: + def dampen_rates(self, system: RodType, time: np.floating) -> None: self.filter_function( - rod.velocity_collection, + system.velocity_collection, self.velocity_filter_term, - rod.omega_collection, + system.omega_collection, self.omega_filter_term, self.filter_order, ) diff --git a/elastica/external_forces.py b/elastica/external_forces.py index e9b5b7238..c3c49549e 100644 --- a/elastica/external_forces.py +++ b/elastica/external_forces.py @@ -274,15 +274,15 @@ def __init__( self.force = (force * direction).reshape(3, 1) def apply_forces( - self, rod: "RodType | RigidBodyType", time: np.floating = np.float64(0.0) + self, system: "RodType | RigidBodyType", time: np.floating = np.float64(0.0) ) -> None: - force_on_one_element = self.force / rod.n_elems + force_on_one_element = self.force / system.n_elems - rod.external_forces += force_on_one_element + system.external_forces += force_on_one_element # Because mass of first and last node is half - rod.external_forces[..., 0] -= 0.5 * force_on_one_element[:, 0] - rod.external_forces[..., -1] -= 0.5 * force_on_one_element[:, 0] + system.external_forces[..., 0] -= 0.5 * force_on_one_element[:, 0] + system.external_forces[..., -1] -= 0.5 * force_on_one_element[:, 0] class MuscleTorques(NoForces): @@ -373,7 +373,7 @@ def __init__( self.my_spline = np.full_like(self.s, fill_value=1.0) def apply_torques( - self, rod: "RodType | RigidBodyType", time: np.floating = np.float64(0.0) + self, system: "RodType | RigidBodyType", time: np.floating = np.float64(0.0) ) -> None: self.compute_muscle_torques( time, @@ -384,8 +384,8 @@ def apply_torques( self.phase_shift, self.ramp_up_time, self.direction, - rod.director_collection, - rod.external_torques, + system.director_collection, + system.external_torques, ) @staticmethod diff --git a/tests/test_dissipation.py b/tests/test_dissipation.py index 49fde7126..8c48c6807 100644 --- a/tests/test_dissipation.py +++ b/tests/test_dissipation.py @@ -151,7 +151,7 @@ def test_laplace_dissipation_filter_for_constant_field(filter_order): ) test_rod.velocity_collection[...] = 2.0 test_rod.omega_collection[...] = 3.0 - filter_damper.dampen_rates(rod=test_rod, time=0) + filter_damper.dampen_rates(system=test_rod, time=0) # filter should keep a spatially invariant field unaffected post_damping_velocity_collection = 2.0 * np.ones_like(test_rod.velocity_collection) post_damping_omega_collection = 3.0 * np.ones_like(test_rod.omega_collection) @@ -178,7 +178,7 @@ def test_laplace_dissipation_filter_for_flip_flop_field(): test_rod.omega_collection[..., 1::2] = 3.0 pre_damping_velocity_collection = test_rod.velocity_collection.copy() pre_damping_omega_collection = test_rod.omega_collection.copy() - filter_damper.dampen_rates(rod=test_rod, time=0) + filter_damper.dampen_rates(system=test_rod, time=0) post_damping_velocity_collection = np.zeros_like(test_rod.velocity_collection) post_damping_omega_collection = np.zeros_like(test_rod.omega_collection) # filter should remove the flip-flop mode th give the average constant mode @@ -243,7 +243,7 @@ def test_laplace_dissipation_filter_for_constant_field_for_ring_rod(filter_order ) test_rod.velocity_collection[...] = 2.0 test_rod.omega_collection[...] = 3.0 - filter_damper.dampen_rates(rod=test_rod, time=0) + filter_damper.dampen_rates(system=test_rod, time=0) # filter should keep a spatially invariant field unaffected post_damping_velocity_collection = 2.0 * np.ones_like(test_rod.velocity_collection) post_damping_omega_collection = 3.0 * np.ones_like(test_rod.omega_collection) @@ -270,7 +270,7 @@ def test_laplace_dissipation_filter_for_flip_flop_field_for_ring_rod(): test_rod.omega_collection[..., 1::2] = 3.0 pre_damping_velocity_collection = test_rod.velocity_collection.copy() pre_damping_omega_collection = test_rod.omega_collection.copy() - filter_damper.dampen_rates(rod=test_rod, time=0) + filter_damper.dampen_rates(system=test_rod, time=0) post_damping_velocity_collection = np.zeros_like(test_rod.velocity_collection) post_damping_omega_collection = np.zeros_like(test_rod.omega_collection) # filter should remove the flip-flop mode th give the average constant mode From 069e2cd0cd2707f2ec44ffd83587c468039058f5 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Tue, 25 Jun 2024 09:48:05 -0500 Subject: [PATCH 120/134] fix: disallow None default value, instead use empty tuple --- elastica/boundary_conditions.py | 4 ++-- elastica/modules/constraints.py | 4 ++-- elastica/typing.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/elastica/boundary_conditions.py b/elastica/boundary_conditions.py index 0a324f119..f71133d52 100644 --- a/elastica/boundary_conditions.py +++ b/elastica/boundary_conditions.py @@ -40,8 +40,8 @@ class ConstraintBase(ABC, Generic[S]): def __init__( self, *args: Any, - constrained_position_idx: ConstrainingIndex, - constrained_director_idx: ConstrainingIndex, + constrained_position_idx: ConstrainingIndex = (), + constrained_director_idx: ConstrainingIndex = (), **kwargs: Any, ) -> None: """Initialize boundary condition""" diff --git a/elastica/modules/constraints.py b/elastica/modules/constraints.py index 2ba907f0c..75ec8e4dd 100644 --- a/elastica/modules/constraints.py +++ b/elastica/modules/constraints.py @@ -151,8 +151,8 @@ def using( self, cls: Type[ConstraintBase], *args: Any, - constrained_position_idx: ConstrainingIndex = None, - constrained_director_idx: ConstrainingIndex = None, + constrained_position_idx: ConstrainingIndex = (), + constrained_director_idx: ConstrainingIndex = (), **kwargs: Any, ) -> Self: """ diff --git a/elastica/typing.py b/elastica/typing.py index aa9928daa..07ac71f56 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -58,7 +58,7 @@ # Indexing types # TODO: Maybe just use slice?? -ConstrainingIndex: TypeAlias = list[int] | tuple[int] | np.typing.NDArray | None +ConstrainingIndex: TypeAlias = list[int] | tuple[int, ...] | np.typing.NDArray ConnectionIndex: TypeAlias = ( int | np.int_ | list[int] | tuple[int] | np.typing.NDArray | None ) From c4584f40fdffe9ac30d0ff5153199f9326893f2d Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Tue, 25 Jun 2024 10:37:03 -0500 Subject: [PATCH 121/134] fix: experimental parallel connection to be batch-wise operation --- .../parallel_connection.py | 9 +-- examples/MuscularSnake/muscular_snake.py | 74 +++++++++---------- 2 files changed, 40 insertions(+), 43 deletions(-) diff --git a/elastica/experimental/connection_contact_joint/parallel_connection.py b/elastica/experimental/connection_contact_joint/parallel_connection.py index 567ab89f4..060af8fc8 100644 --- a/elastica/experimental/connection_contact_joint/parallel_connection.py +++ b/elastica/experimental/connection_contact_joint/parallel_connection.py @@ -90,10 +90,10 @@ def __init__( self.rod_one_direction_vec_in_material_frame = np.array( rod_one_direction_vec_in_material_frame - ).T + ) self.rod_two_direction_vec_in_material_frame = np.array( rod_two_direction_vec_in_material_frame - ).T + ) # Apply force is same as free joint def apply_forces(self, system_one, index_one, system_two, index_two): @@ -150,7 +150,6 @@ def _apply_forces( rod_one_external_forces, rod_two_external_forces, ): - rod_one_to_rod_two_connection_vec = _batch_matvec( _batch_matrix_transpose(rod_one_director_collection[:, :, index_one]), rod_one_direction_vec_in_material_frame, @@ -286,8 +285,8 @@ def _apply_torques( rod_two_external_torques, ): # Compute torques due to the connection forces - torque_on_rod_one = np.cross(rod_one_rd2, spring_force) - torque_on_rod_two = np.cross(rod_two_rd2, -spring_force) + torque_on_rod_one = _batch_cross(rod_one_rd2, spring_force) + torque_on_rod_two = _batch_cross(rod_two_rd2, -spring_force) torque_on_rod_one_material_frame = _batch_matvec( rod_one_director_collection[:, :, index_one], torque_on_rod_one diff --git a/examples/MuscularSnake/muscular_snake.py b/examples/MuscularSnake/muscular_snake.py index c2e1c7bb6..2d224c63f 100644 --- a/examples/MuscularSnake/muscular_snake.py +++ b/examples/MuscularSnake/muscular_snake.py @@ -290,45 +290,43 @@ class MuscularSnakeSimulator( offset_btw_rods.copy(), ] ) - for k in range(rod_two.n_elems): - rod_one_index = k + muscle_start_connection_index[idx] - rod_two_index = k - k_conn = ( - rod_one.radius[rod_one_index] - * rod_two.radius[rod_two_index] - / (rod_one.radius[rod_one_index] + rod_two.radius[rod_two_index]) - * body_elem_length - * E - / (rod_one.radius[rod_one_index] + rod_two.radius[rod_two_index]) - ) - if k < 12 or k >= 27: - scale = 1 * 2 - scale_contact = 20 - else: - scale = 0.01 * 5 - scale_contact = 20 - - muscular_snake_simulator.connect( - first_rod=rod_one, - second_rod=rod_two, - first_connect_idx=rod_one_index, - second_connect_idx=rod_two_index, - ).using( - SurfaceJointSideBySide, - k=k_conn * scale, - nu=1e-4, - k_repulsive=k_conn * scale_contact, - rod_one_direction_vec_in_material_frame=rod_one_direction_vec_in_material_frame[ - ..., k - ], - rod_two_direction_vec_in_material_frame=rod_two_direction_vec_in_material_frame[ - ..., k - ], - offset_btw_rods=offset_btw_rods[k], - post_processing_dict=straight_straight_rod_connection_post_processing_dict, - step_skip=step_skip, - ) + ks = np.arange(rod_two.n_elems) + scale = np.ones(rod_two.n_elems) * 1 * 2 + scale[12:27] = 0.01 * 5 + scale_contact = np.ones(rod_two.n_elems) * 20 + scale_contact[12:27] = 20 + rod_one_index = ks + muscle_start_connection_index[idx] + rod_two_index = ks + k_conn = ( + rod_one.radius[rod_one_index] + * rod_two.radius[rod_two_index] + / (rod_one.radius[rod_one_index] + rod_two.radius[rod_two_index]) + * body_elem_length + * E + / (rod_one.radius[rod_one_index] + rod_two.radius[rod_two_index]) + ) + + muscular_snake_simulator.connect( + first_rod=rod_one, + second_rod=rod_two, + first_connect_idx=rod_one_index, + second_connect_idx=rod_two_index, + ).using( + SurfaceJointSideBySide, + k=k_conn * scale, + nu=1e-4, + k_repulsive=k_conn * scale_contact, + rod_one_direction_vec_in_material_frame=rod_one_direction_vec_in_material_frame[ + ..., ks + ], + rod_two_direction_vec_in_material_frame=rod_two_direction_vec_in_material_frame[ + ..., ks + ], + offset_btw_rods=offset_btw_rods[ks], + post_processing_dict=straight_straight_rod_connection_post_processing_dict, + step_skip=step_skip, + ) # Friction forces # Only apply to the snake body. From 9b267aff9722a97b416aee6e88e5eae245ca8e64 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Tue, 25 Jun 2024 10:39:42 -0500 Subject: [PATCH 122/134] check example cases, fix outdated syntax --- elastica/_linalg.py | 2 ++ elastica/boundary_conditions.py | 8 ++------ elastica/modules/base_system.py | 4 ++++ .../ContinuumSnakeWithLiftingWaveCase/snake_contact.py | 1 + .../rigid_cylinder_rotational_motion_case.py | 2 +- .../rigid_cylinder_translational_motion_case.py | 2 +- .../RigidBodyCases/rigid_sphere_rotational_motion_case.py | 2 +- .../rigid_sphere_translational_motion_case.py | 2 +- 8 files changed, 13 insertions(+), 10 deletions(-) diff --git a/elastica/_linalg.py b/elastica/_linalg.py index 5097c7093..000af5b5d 100644 --- a/elastica/_linalg.py +++ b/elastica/_linalg.py @@ -376,6 +376,8 @@ def _batch_vector_sum( def _batch_matrix_transpose(input_matrix: NDArray[np.floating]) -> NDArray[np.floating]: """ This function takes an batch input matrix and transpose it. + [i,j,k] -> [j,i,k] + Parameters ---------- input_matrix diff --git a/elastica/boundary_conditions.py b/elastica/boundary_conditions.py index f71133d52..9b5339937 100644 --- a/elastica/boundary_conditions.py +++ b/elastica/boundary_conditions.py @@ -66,18 +66,15 @@ def system(self) -> S: @property def constrained_position_idx(self) -> np.ndarray: """get position-indices passed to "using" """ - # TODO: This should be immutable somehow return self._constrained_position_idx @property def constrained_director_idx(self) -> np.ndarray: """get director-indices passed to "using" """ - # TODO: This should be immutable somehow return self._constrained_director_idx @abstractmethod def constrain_values(self, system: S, time: np.floating) -> None: - # TODO: In the future, we can remove rod and use self.system """ Constrain values (position and/or directors) of a rod object. @@ -88,11 +85,10 @@ def constrain_values(self, system: S, time: np.floating) -> None: time : float The time of simulation. """ - raise NotImplementedError + pass @abstractmethod def constrain_rates(self, system: S, time: np.floating) -> None: - # TODO: In the future, we can remove rod and use self.system """ Constrain rates (velocity and/or omega) of a rod object. @@ -104,7 +100,7 @@ def constrain_rates(self, system: S, time: np.floating) -> None: The time of simulation. """ - raise NotImplementedError + pass class FreeBC(ConstraintBase): diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index ff32cf922..faf922a1a 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -186,6 +186,10 @@ def finalize(self) -> None: # construct memory block self.__final_systems = construct_memory_block_structures(self._systems) + self._systems.extend( + self.__final_systems + ) # FIXME: We need this to make ring-rod working. + # But probably need to be refactored # TODO: try to remove the _systems list for memory optimization # self._systems.clear() # del self._systems diff --git a/examples/ContinuumSnakeWithLiftingWaveCase/snake_contact.py b/examples/ContinuumSnakeWithLiftingWaveCase/snake_contact.py index c27ffd37c..464b1fd36 100755 --- a/examples/ContinuumSnakeWithLiftingWaveCase/snake_contact.py +++ b/examples/ContinuumSnakeWithLiftingWaveCase/snake_contact.py @@ -1,4 +1,5 @@ __doc__ = """Rod plane contact with anistropic friction (no static friction)""" +from typing import Type import numpy as np from elastica._linalg import ( diff --git a/examples/RigidBodyCases/rigid_cylinder_rotational_motion_case.py b/examples/RigidBodyCases/rigid_cylinder_rotational_motion_case.py index caa529fb3..f33b6ccd0 100644 --- a/examples/RigidBodyCases/rigid_cylinder_rotational_motion_case.py +++ b/examples/RigidBodyCases/rigid_cylinder_rotational_motion_case.py @@ -52,7 +52,7 @@ def __init__(self, torque, direction=np.array([0.0, 0.0, 0.0])): super(PointCoupleToCenter, self).__init__() self.torque = (torque * direction).reshape(3, 1) - def apply_forces(self, system, time: np.float = 0.0): + def apply_forces(self, system, time: np.floating = np.float64(0.0)): system.external_torques += np.einsum( "ijk, jk->ik", system.director_collection, self.torque ) diff --git a/examples/RigidBodyCases/rigid_cylinder_translational_motion_case.py b/examples/RigidBodyCases/rigid_cylinder_translational_motion_case.py index 0f95ffcb0..0e21bc0ba 100644 --- a/examples/RigidBodyCases/rigid_cylinder_translational_motion_case.py +++ b/examples/RigidBodyCases/rigid_cylinder_translational_motion_case.py @@ -52,7 +52,7 @@ def __init__(self, force, direction=np.array([0.0, 0.0, 0.0])): super(PointForceToCenter, self).__init__() self.force = (force * direction).reshape(3, 1) - def apply_forces(self, system, time: np.float = 0.0): + def apply_forces(self, system, time: np.floating = np.float64(0.0)): system.external_forces += self.force # Add point force on the rod diff --git a/examples/RigidBodyCases/rigid_sphere_rotational_motion_case.py b/examples/RigidBodyCases/rigid_sphere_rotational_motion_case.py index 5c2f5d837..2f377209a 100644 --- a/examples/RigidBodyCases/rigid_sphere_rotational_motion_case.py +++ b/examples/RigidBodyCases/rigid_sphere_rotational_motion_case.py @@ -43,7 +43,7 @@ def __init__(self, torque, direction=np.array([0.0, 0.0, 0.0])): super(PointCoupleToCenter, self).__init__() self.torque = (torque * direction).reshape(3, 1) - def apply_forces(self, system, time: np.float = 0.0): + def apply_forces(self, system, time: np.floating = np.float64(0.0)): system.external_torques += np.einsum( "ijk, jk->ik", system.director_collection, self.torque ) diff --git a/examples/RigidBodyCases/rigid_sphere_translational_motion_case.py b/examples/RigidBodyCases/rigid_sphere_translational_motion_case.py index 3a3ab0719..3ececd33e 100644 --- a/examples/RigidBodyCases/rigid_sphere_translational_motion_case.py +++ b/examples/RigidBodyCases/rigid_sphere_translational_motion_case.py @@ -44,7 +44,7 @@ def __init__(self, force, direction=np.array([0.0, 0.0, 0.0])): super(PointForceToCenter, self).__init__() self.force = (force * direction).reshape(3, 1) - def apply_forces(self, system, time: np.float = 0.0): + def apply_forces(self, system, time: np.floating = np.float64(0.0)): system.external_forces += self.force # Add point force on the rod From 9d64f193f6150c1556376d894b37c7f84529fbfe Mon Sep 17 00:00:00 2001 From: Arman Tekinalp <53585636+armantekinalp@users.noreply.github.com> Date: Tue, 25 Jun 2024 17:04:00 -0500 Subject: [PATCH 123/134] Update elastica/rod/cosserat_rod.py --- elastica/rod/cosserat_rod.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/elastica/rod/cosserat_rod.py b/elastica/rod/cosserat_rod.py index ea79854a1..b309dbf46 100644 --- a/elastica/rod/cosserat_rod.py +++ b/elastica/rod/cosserat_rod.py @@ -188,7 +188,7 @@ def __init__( internal_couple: NDArray[np.floating], ring_rod_flag: bool, ) -> None: - self.n_nodes = n_elements + 1 if ring_rod_flag else n_elements + self.n_nodes = n_elements + 1 if not ring_rod_flag else n_elements self.n_elems = n_elements self.position_collection = position self.velocity_collection = velocity From 2ff4e4d1ec96da7ad69665e7046afad33bee7106 Mon Sep 17 00:00:00 2001 From: Arman Tekinalp <53585636+armantekinalp@users.noreply.github.com> Date: Tue, 25 Jun 2024 17:08:12 -0500 Subject: [PATCH 124/134] Update elastica/external_forces.py --- elastica/external_forces.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/elastica/external_forces.py b/elastica/external_forces.py index c3c49549e..fc46687a7 100644 --- a/elastica/external_forces.py +++ b/elastica/external_forces.py @@ -508,8 +508,8 @@ def __init__( start_force_mag: float, end_force_mag: float, ramp_up_time: float = 0.0, - tangent_direction: NDArray[np.floating] = np.array([0, 0, 1]), - normal_direction: NDArray[np.floating] = np.array([0, 1, 0]), + tangent_direction: NDArray[np.floating] = np.array([0., 0., 1.]), + normal_direction: NDArray[np.floating] = np.array([0., 1., 0.]), ) -> None: """ From be228f0e09f188891d84dd83c00716b2ce51dabe Mon Sep 17 00:00:00 2001 From: Arman Tekinalp <53585636+armantekinalp@users.noreply.github.com> Date: Tue, 25 Jun 2024 17:09:35 -0500 Subject: [PATCH 125/134] Update elastica/modules/protocol.py --- elastica/modules/protocol.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index 5de97c7e2..eba0c4c33 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -1,5 +1,5 @@ from typing import Protocol, Generator, TypeVar, Any, Type, overload -from typing_extensions import Self # 3.11: from typing import Self +from typing_extensions import Self # python 3.11: from typing import Self from abc import abstractmethod From bb4b2c5e8309e425a0566055feae9ea917f6e1ce Mon Sep 17 00:00:00 2001 From: Arman Tekinalp <53585636+armantekinalp@users.noreply.github.com> Date: Tue, 25 Jun 2024 17:10:23 -0500 Subject: [PATCH 126/134] Update elastica/typing.py --- elastica/typing.py | 1 - 1 file changed, 1 deletion(-) diff --git a/elastica/typing.py b/elastica/typing.py index 07ac71f56..1ba1a4ee3 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -40,7 +40,6 @@ SystemIdxType: TypeAlias = int BlockSystemType: TypeAlias = "BlockSystemProtocol" -# NoForces | NoContact | FreeJoint | FreeBC | DamperBase | CallBackBaseClass # Mostly used in explicit stepper: for symplectic, use kinetic and dynamic state StateType: TypeAlias = "State" From e3ac13eeb889de40c82ae7675c9bdc300ccbdc8e Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Tue, 25 Jun 2024 22:14:15 -0500 Subject: [PATCH 127/134] remove ambiguity: np.integer -> np.int32, np.floating -> np.float64 --- elastica/_calculus.py | 24 +- elastica/_contact_functions.py | 210 +++++------ elastica/_linalg.py | 46 +-- elastica/_rotations.py | 26 +- elastica/_synchronize_periodic_boundary.py | 10 +- elastica/boundary_conditions.py | 86 ++--- elastica/callback_functions.py | 8 +- elastica/collision/AABBCollection.py | 20 +- elastica/contact_forces.py | 6 +- elastica/contact_utils.py | 78 ++-- elastica/dissipation.py | 26 +- elastica/external_forces.py | 84 ++--- elastica/interaction.py | 46 ++- elastica/joint.py | 6 +- .../memory_block/memory_block_rigid_body.py | 2 +- elastica/memory_block/memory_block_rod.py | 8 +- elastica/memory_block/utils.py | 68 ++-- elastica/mesh/protocol.py | 6 +- elastica/modules/base_system.py | 8 +- elastica/modules/callbacks.py | 2 +- elastica/modules/connections.py | 2 +- elastica/modules/constraints.py | 4 +- elastica/modules/contact.py | 2 +- elastica/modules/damping.py | 2 +- elastica/modules/protocol.py | 16 +- .../_reset_ghost_vector_or_scalar.py | 24 +- elastica/restart.py | 2 +- elastica/rigidbody/cylinder.py | 14 +- elastica/rigidbody/mesh_rigid_body.py | 8 +- elastica/rigidbody/protocol.py | 14 +- elastica/rigidbody/rigid_body.py | 46 +-- elastica/rigidbody/sphere.py | 4 +- elastica/rod/cosserat_rod.py | 354 +++++++++--------- elastica/rod/data_structures.py | 72 ++-- elastica/rod/factory_function.py | 84 ++--- elastica/rod/knot_theory.py | 60 +-- elastica/rod/protocol.py | 58 +-- elastica/rod/rod_base.py | 30 +- elastica/surface/plane.py | 2 +- elastica/surface/surface_base.py | 4 +- elastica/systems/protocol.py | 40 +- elastica/timestepper/__init__.py | 2 +- elastica/timestepper/explicit_steppers.py | 68 ++-- elastica/timestepper/protocol.py | 8 +- elastica/timestepper/symplectic_steppers.py | 40 +- elastica/transformations.py | 18 +- elastica/utils.py | 2 +- examples/MuscularSnake/muscular_snake.py | 10 +- .../rigid_cylinder_rotational_motion_case.py | 2 +- ...igid_cylinder_translational_motion_case.py | 2 +- .../rigid_sphere_rotational_motion_case.py | 2 +- .../rigid_sphere_translational_motion_case.py | 2 +- tests/test_math/test_governing_equations.py | 8 +- tests/test_memory_block_validity.py | 72 ++-- tests/test_modules/test_memory_block_utils.py | 54 +-- tests/test_synchronize_periodic_boundary.py | 6 +- 56 files changed, 949 insertions(+), 959 deletions(-) diff --git a/elastica/_calculus.py b/elastica/_calculus.py index 1639eafd2..c6704480b 100644 --- a/elastica/_calculus.py +++ b/elastica/_calculus.py @@ -12,7 +12,7 @@ @functools.lru_cache(maxsize=2) -def _get_zero_array(dim: int, ndim: int) -> Union[float, NDArray[np.floating], None]: +def _get_zero_array(dim: int, ndim: int) -> Union[float, NDArray[np.float64], None]: if ndim == 1: return 0.0 if ndim == 2: @@ -22,7 +22,7 @@ def _get_zero_array(dim: int, ndim: int) -> Union[float, NDArray[np.floating], N @njit(cache=True) # type: ignore -def _trapezoidal(array_collection: NDArray[np.floating]) -> NDArray[np.floating]: +def _trapezoidal(array_collection: NDArray[np.float64]) -> NDArray[np.float64]: """ Simple trapezoidal quadrature rule with zero at end-points, in a dimension agnostic way @@ -69,8 +69,8 @@ def _trapezoidal(array_collection: NDArray[np.floating]) -> NDArray[np.floating] @njit(cache=True) # type: ignore def _trapezoidal_for_block_structure( - array_collection: NDArray[np.floating], ghost_idx: NDArray[np.integer] -) -> NDArray[np.floating]: + array_collection: NDArray[np.float64], ghost_idx: NDArray[np.int32] +) -> NDArray[np.float64]: """ Simple trapezoidal quadrature rule with zero at end-points, in a dimension agnostic way. This form specifically for the block structure implementation and there is a reset function call, to reset @@ -123,8 +123,8 @@ def _trapezoidal_for_block_structure( @njit(cache=True) # type: ignore def _two_point_difference( - array_collection: NDArray[np.floating], -) -> NDArray[np.floating]: + array_collection: NDArray[np.float64], +) -> NDArray[np.float64]: """ This function does differentiation. @@ -166,8 +166,8 @@ def _two_point_difference( @njit(cache=True) # type: ignore def _two_point_difference_for_block_structure( - array_collection: NDArray[np.floating], ghost_idx: NDArray[np.integer] -) -> NDArray[np.floating]: + array_collection: NDArray[np.float64], ghost_idx: NDArray[np.int32] +) -> NDArray[np.float64]: """ This function does the differentiation, for Cosserat rod model equations. This form specifically for the block structure implementation and there is a reset function call, to @@ -218,7 +218,7 @@ def _two_point_difference_for_block_structure( @njit(cache=True) # type: ignore -def _difference(vector: NDArray[np.floating]) -> NDArray[np.floating]: +def _difference(vector: NDArray[np.float64]) -> NDArray[np.float64]: """ This function computes difference between elements of a batch vector. @@ -249,7 +249,7 @@ def _difference(vector: NDArray[np.floating]) -> NDArray[np.floating]: @njit(cache=True) # type: ignore -def _average(vector: NDArray[np.floating]) -> NDArray[np.floating]: +def _average(vector: NDArray[np.float64]) -> NDArray[np.float64]: """ This function computes the average between elements of a vector. @@ -280,8 +280,8 @@ def _average(vector: NDArray[np.floating]) -> NDArray[np.floating]: @njit(cache=True) # type: ignore def _clip_array( - input_array: NDArray[np.floating], vmin: np.floating, vmax: np.floating -) -> NDArray[np.floating]: + input_array: NDArray[np.float64], vmin: np.float64, vmax: np.float64 +) -> NDArray[np.float64]: """ This function clips an array values between user defined minimum and maximum diff --git a/elastica/_contact_functions.py b/elastica/_contact_functions.py index be33436ac..d1456092f 100644 --- a/elastica/_contact_functions.py +++ b/elastica/_contact_functions.py @@ -30,24 +30,24 @@ @njit(cache=True) # type: ignore def _calculate_contact_forces_rod_cylinder( - x_collection_rod: NDArray[np.floating], - edge_collection_rod: NDArray[np.floating], - x_cylinder_center: NDArray[np.floating], - x_cylinder_tip: NDArray[np.floating], - edge_cylinder: NDArray[np.floating], - radii_sum: NDArray[np.floating], - length_sum: NDArray[np.floating], - internal_forces_rod: NDArray[np.floating], - external_forces_rod: NDArray[np.floating], - external_forces_cylinder: NDArray[np.floating], - external_torques_cylinder: NDArray[np.floating], - cylinder_director_collection: NDArray[np.floating], - velocity_rod: NDArray[np.floating], - velocity_cylinder: NDArray[np.floating], - contact_k: np.floating, - contact_nu: np.floating, - velocity_damping_coefficient: np.floating, - friction_coefficient: np.floating, + x_collection_rod: NDArray[np.float64], + edge_collection_rod: NDArray[np.float64], + x_cylinder_center: NDArray[np.float64], + x_cylinder_tip: NDArray[np.float64], + edge_cylinder: NDArray[np.float64], + radii_sum: NDArray[np.float64], + length_sum: NDArray[np.float64], + internal_forces_rod: NDArray[np.float64], + external_forces_rod: NDArray[np.float64], + external_forces_cylinder: NDArray[np.float64], + external_torques_cylinder: NDArray[np.float64], + cylinder_director_collection: NDArray[np.float64], + velocity_rod: NDArray[np.float64], + velocity_cylinder: NDArray[np.float64], + contact_k: np.float64, + contact_nu: np.float64, + velocity_damping_coefficient: np.float64, + friction_coefficient: np.float64, ) -> None: # We already pass in only the first n_elem x n_points = x_collection_rod.shape[1] @@ -157,22 +157,22 @@ def _calculate_contact_forces_rod_cylinder( @njit(cache=True) # type: ignore def _calculate_contact_forces_rod_rod( - x_collection_rod_one: NDArray[np.floating], - radius_rod_one: NDArray[np.floating], - length_rod_one: NDArray[np.floating], - tangent_rod_one: NDArray[np.floating], - velocity_rod_one: NDArray[np.floating], - internal_forces_rod_one: NDArray[np.floating], - external_forces_rod_one: NDArray[np.floating], - x_collection_rod_two: NDArray[np.floating], - radius_rod_two: NDArray[np.floating], - length_rod_two: NDArray[np.floating], - tangent_rod_two: NDArray[np.floating], - velocity_rod_two: NDArray[np.floating], - internal_forces_rod_two: NDArray[np.floating], - external_forces_rod_two: NDArray[np.floating], - contact_k: np.floating, - contact_nu: np.floating, + x_collection_rod_one: NDArray[np.float64], + radius_rod_one: NDArray[np.float64], + length_rod_one: NDArray[np.float64], + tangent_rod_one: NDArray[np.float64], + velocity_rod_one: NDArray[np.float64], + internal_forces_rod_one: NDArray[np.float64], + external_forces_rod_one: NDArray[np.float64], + x_collection_rod_two: NDArray[np.float64], + radius_rod_two: NDArray[np.float64], + length_rod_two: NDArray[np.float64], + tangent_rod_two: NDArray[np.float64], + velocity_rod_two: NDArray[np.float64], + internal_forces_rod_two: NDArray[np.float64], + external_forces_rod_two: NDArray[np.float64], + contact_k: np.float64, + contact_nu: np.float64, ) -> None: # We already pass in only the first n_elem x n_points_rod_one = x_collection_rod_one.shape[1] @@ -274,14 +274,14 @@ def _calculate_contact_forces_rod_rod( @njit(cache=True) # type: ignore def _calculate_contact_forces_self_rod( - x_collection_rod: NDArray[np.floating], - radius_rod: NDArray[np.floating], - length_rod: NDArray[np.floating], - tangent_rod: NDArray[np.floating], - velocity_rod: NDArray[np.floating], - external_forces_rod: NDArray[np.floating], - contact_k: np.floating, - contact_nu: np.floating, + x_collection_rod: NDArray[np.float64], + radius_rod: NDArray[np.float64], + length_rod: NDArray[np.float64], + tangent_rod: NDArray[np.float64], + velocity_rod: NDArray[np.float64], + external_forces_rod: NDArray[np.float64], + contact_k: np.float64, + contact_nu: np.float64, ) -> None: # We already pass in only the first n_elem x n_points_rod = x_collection_rod.shape[1] @@ -362,24 +362,24 @@ def _calculate_contact_forces_self_rod( @njit(cache=True) # type: ignore def _calculate_contact_forces_rod_sphere( - x_collection_rod: NDArray[np.floating], - edge_collection_rod: NDArray[np.floating], - x_sphere_center: NDArray[np.floating], - x_sphere_tip: NDArray[np.floating], - edge_sphere: NDArray[np.floating], - radii_sum: NDArray[np.floating], - length_sum: NDArray[np.floating], - internal_forces_rod: NDArray[np.floating], - external_forces_rod: NDArray[np.floating], - external_forces_sphere: NDArray[np.floating], - external_torques_sphere: NDArray[np.floating], - sphere_director_collection: NDArray[np.floating], - velocity_rod: NDArray[np.floating], - velocity_sphere: NDArray[np.floating], - contact_k: np.floating, - contact_nu: np.floating, - velocity_damping_coefficient: np.floating, - friction_coefficient: np.floating, + x_collection_rod: NDArray[np.float64], + edge_collection_rod: NDArray[np.float64], + x_sphere_center: NDArray[np.float64], + x_sphere_tip: NDArray[np.float64], + edge_sphere: NDArray[np.float64], + radii_sum: NDArray[np.float64], + length_sum: NDArray[np.float64], + internal_forces_rod: NDArray[np.float64], + external_forces_rod: NDArray[np.float64], + external_forces_sphere: NDArray[np.float64], + external_torques_sphere: NDArray[np.float64], + sphere_director_collection: NDArray[np.float64], + velocity_rod: NDArray[np.float64], + velocity_sphere: NDArray[np.float64], + contact_k: np.float64, + contact_nu: np.float64, + velocity_damping_coefficient: np.float64, + friction_coefficient: np.float64, ) -> None: # We already pass in only the first n_elem x n_points = x_collection_rod.shape[1] @@ -488,18 +488,18 @@ def _calculate_contact_forces_rod_sphere( @njit(cache=True) # type: ignore def _calculate_contact_forces_rod_plane( - plane_origin: NDArray[np.floating], - plane_normal: NDArray[np.floating], - surface_tol: np.floating, - k: np.floating, - nu: np.floating, - radius: NDArray[np.floating], - mass: NDArray[np.floating], - position_collection: NDArray[np.floating], - velocity_collection: NDArray[np.floating], - internal_forces: NDArray[np.floating], - external_forces: NDArray[np.floating], -) -> tuple[NDArray[np.floating], NDArray[np.intp]]: + plane_origin: NDArray[np.float64], + plane_normal: NDArray[np.float64], + surface_tol: np.float64, + k: np.float64, + nu: np.float64, + radius: NDArray[np.float64], + mass: NDArray[np.float64], + position_collection: NDArray[np.float64], + velocity_collection: NDArray[np.float64], + internal_forces: NDArray[np.float64], + external_forces: NDArray[np.float64], +) -> tuple[NDArray[np.float64], NDArray[np.intp]]: """ This function computes the plane force response on the element, in the case of contact. Contact model given in Eqn 4.8 Gazzola et. al. RSoS 2018 paper @@ -573,29 +573,29 @@ def _calculate_contact_forces_rod_plane( @njit(cache=True) # type: ignore def _calculate_contact_forces_rod_plane_with_anisotropic_friction( - plane_origin: NDArray[np.floating], - plane_normal: NDArray[np.floating], - surface_tol: np.floating, - slip_velocity_tol: np.floating, - k: np.floating, - nu: np.floating, - kinetic_mu_forward: np.floating, - kinetic_mu_backward: np.floating, - kinetic_mu_sideways: np.floating, - static_mu_forward: np.floating, - static_mu_backward: np.floating, - static_mu_sideways: np.floating, - radius: NDArray[np.floating], - mass: NDArray[np.floating], - tangents: NDArray[np.floating], - position_collection: NDArray[np.floating], - director_collection: NDArray[np.floating], - velocity_collection: NDArray[np.floating], - omega_collection: NDArray[np.floating], - internal_forces: NDArray[np.floating], - external_forces: NDArray[np.floating], - internal_torques: NDArray[np.floating], - external_torques: NDArray[np.floating], + plane_origin: NDArray[np.float64], + plane_normal: NDArray[np.float64], + surface_tol: np.float64, + slip_velocity_tol: np.float64, + k: np.float64, + nu: np.float64, + kinetic_mu_forward: np.float64, + kinetic_mu_backward: np.float64, + kinetic_mu_sideways: np.float64, + static_mu_forward: np.float64, + static_mu_backward: np.float64, + static_mu_sideways: np.float64, + radius: NDArray[np.float64], + mass: NDArray[np.float64], + tangents: NDArray[np.float64], + position_collection: NDArray[np.float64], + director_collection: NDArray[np.float64], + velocity_collection: NDArray[np.float64], + omega_collection: NDArray[np.float64], + internal_forces: NDArray[np.float64], + external_forces: NDArray[np.float64], + internal_torques: NDArray[np.float64], + external_torques: NDArray[np.float64], ) -> None: ( plane_response_force_mag, @@ -786,16 +786,16 @@ def _calculate_contact_forces_rod_plane_with_anisotropic_friction( @njit(cache=True) # type: ignore def _calculate_contact_forces_cylinder_plane( - plane_origin: NDArray[np.floating], - plane_normal: NDArray[np.floating], - surface_tol: np.floating, - k: np.floating, - nu: np.floating, - length: NDArray[np.floating], - position_collection: NDArray[np.floating], - velocity_collection: NDArray[np.floating], - external_forces: NDArray[np.floating], -) -> tuple[NDArray[np.floating], NDArray[np.intp]]: + plane_origin: NDArray[np.float64], + plane_normal: NDArray[np.float64], + surface_tol: np.float64, + k: np.float64, + nu: np.float64, + length: NDArray[np.float64], + position_collection: NDArray[np.float64], + velocity_collection: NDArray[np.float64], + external_forces: NDArray[np.float64], +) -> tuple[NDArray[np.float64], NDArray[np.intp]]: # Compute plane response force # total_forces = system.internal_forces + system.external_forces diff --git a/elastica/_linalg.py b/elastica/_linalg.py index 000af5b5d..c4ce37d33 100644 --- a/elastica/_linalg.py +++ b/elastica/_linalg.py @@ -9,7 +9,7 @@ @functools.lru_cache(maxsize=1) -def levi_civita_tensor(dim: int) -> NDArray[np.floating]: +def levi_civita_tensor(dim: int) -> NDArray[np.float64]: """ Parameters @@ -30,8 +30,8 @@ def levi_civita_tensor(dim: int) -> NDArray[np.floating]: @njit(cache=True) # type: ignore def _batch_matvec( - matrix_collection: NDArray[np.floating], vector_collection: NDArray[np.floating] -) -> NDArray[np.floating]: + matrix_collection: NDArray[np.float64], vector_collection: NDArray[np.float64] +) -> NDArray[np.float64]: """ This function does batch matrix and batch vector product @@ -63,9 +63,9 @@ def _batch_matvec( @njit(cache=True) # type: ignore def _batch_matmul( - first_matrix_collection: NDArray[np.floating], - second_matrix_collection: NDArray[np.floating], -) -> NDArray[np.floating]: + first_matrix_collection: NDArray[np.float64], + second_matrix_collection: NDArray[np.float64], +) -> NDArray[np.float64]: """ This is batch matrix matrix multiplication function. Only batch of 3x3 matrices can be multiplied. @@ -100,9 +100,9 @@ def _batch_matmul( @njit(cache=True) # type: ignore def _batch_cross( - first_vector_collection: NDArray[np.floating], - second_vector_collection: NDArray[np.floating], -) -> NDArray[np.floating]: + first_vector_collection: NDArray[np.float64], + second_vector_collection: NDArray[np.float64], +) -> NDArray[np.float64]: """ This function does cross product between two batch vectors. @@ -143,8 +143,8 @@ def _batch_cross( @njit(cache=True) # type: ignore def _batch_vec_oneD_vec_cross( - first_vector_collection: NDArray[np.floating], second_vector: NDArray[np.floating] -) -> NDArray[np.floating]: + first_vector_collection: NDArray[np.float64], second_vector: NDArray[np.float64] +) -> NDArray[np.float64]: """ This function does cross product between batch vector and a 1D vector. Idea of having this function is that, for friction calculations, we dont @@ -189,8 +189,8 @@ def _batch_vec_oneD_vec_cross( @njit(cache=True) # type: ignore def _batch_dot( - first_vector: NDArray[np.floating], second_vector: NDArray[np.floating] -) -> NDArray[np.floating]: + first_vector: NDArray[np.float64], second_vector: NDArray[np.float64] +) -> NDArray[np.float64]: """ This function does batch vec and batch vec dot product. Parameters @@ -217,7 +217,7 @@ def _batch_dot( @njit(cache=True) # type: ignore -def _batch_norm(vector: NDArray[np.floating]) -> NDArray[np.floating]: +def _batch_norm(vector: NDArray[np.float64]) -> NDArray[np.float64]: """ This function computes norm of a batch vector Parameters @@ -247,8 +247,8 @@ def _batch_norm(vector: NDArray[np.floating]) -> NDArray[np.floating]: @njit(cache=True) # type: ignore def _batch_product_i_k_to_ik( - vector1: NDArray[np.floating], vector2: NDArray[np.floating] -) -> NDArray[np.floating]: + vector1: NDArray[np.float64], vector2: NDArray[np.float64] +) -> NDArray[np.float64]: """ This function does outer product following 'i,k->ik'. vector1 has shape of 3 and vector 2 has shape of blocksize @@ -278,8 +278,8 @@ def _batch_product_i_k_to_ik( @njit(cache=True) # type: ignore def _batch_product_i_ik_to_k( - vector1: NDArray[np.floating], vector2: NDArray[np.floating] -) -> NDArray[np.floating]: + vector1: NDArray[np.float64], vector2: NDArray[np.float64] +) -> NDArray[np.float64]: """ This function does the following product 'i,ik->k' This function do dot product between a vector of 3 elements @@ -311,8 +311,8 @@ def _batch_product_i_ik_to_k( @njit(cache=True) # type: ignore def _batch_product_k_ik_to_ik( - vector1: NDArray[np.floating], vector2: NDArray[np.floating] -) -> NDArray[np.floating]: + vector1: NDArray[np.float64], vector2: NDArray[np.float64] +) -> NDArray[np.float64]: """ This function does the following product 'k, ik->ik' Parameters @@ -342,8 +342,8 @@ def _batch_product_k_ik_to_ik( @njit(cache=True) # type: ignore def _batch_vector_sum( - vector1: NDArray[np.floating], vector2: NDArray[np.floating] -) -> NDArray[np.floating]: + vector1: NDArray[np.float64], vector2: NDArray[np.float64] +) -> NDArray[np.float64]: """ This function is for summing up two vectors. Although this function is not faster than pure python implementation @@ -373,7 +373,7 @@ def _batch_vector_sum( @njit(cache=True) # type: ignore -def _batch_matrix_transpose(input_matrix: NDArray[np.floating]) -> NDArray[np.floating]: +def _batch_matrix_transpose(input_matrix: NDArray[np.float64]) -> NDArray[np.float64]: """ This function takes an batch input matrix and transpose it. [i,j,k] -> [j,i,k] diff --git a/elastica/_rotations.py b/elastica/_rotations.py index b7d1180a6..11d279eee 100644 --- a/elastica/_rotations.py +++ b/elastica/_rotations.py @@ -17,8 +17,8 @@ @njit(cache=True) # type: ignore def _get_rotation_matrix( - scale: np.floating, axis_collection: NDArray[np.floating] -) -> NDArray[np.floating]: + scale: np.float64, axis_collection: NDArray[np.float64] +) -> NDArray[np.float64]: blocksize = axis_collection.shape[1] rot_mat = np.empty((3, 3, blocksize)) @@ -53,10 +53,10 @@ def _get_rotation_matrix( @njit(cache=True) # type: ignore def _rotate( - director_collection: NDArray[np.floating], - scale: np.floating, - axis_collection: NDArray[np.floating], -) -> NDArray[np.floating]: + director_collection: NDArray[np.float64], + scale: np.float64, + axis_collection: NDArray[np.float64], +) -> NDArray[np.float64]: """ Does alibi rotations https://en.wikipedia.org/wiki/Rotation_matrix#Ambiguities @@ -81,7 +81,7 @@ def _rotate( @njit(cache=True) # type: ignore -def _inv_rotate(director_collection: NDArray[np.floating]) -> NDArray[np.floating]: +def _inv_rotate(director_collection: NDArray[np.float64]) -> NDArray[np.float64]: """ Calculated rate of change using Rodrigues' formula @@ -244,7 +244,7 @@ def _get_diag_map(dim: int) -> tuple[int, ...]: return tuple([dim_iter * (dim + 1) for dim_iter in range(dim)]) -def _skew_symmetrize(vector: NDArray[np.floating]) -> NDArray[np.floating]: +def _skew_symmetrize(vector: NDArray[np.float64]) -> NDArray[np.float64]: """ Parameters @@ -279,7 +279,7 @@ def _skew_symmetrize(vector: NDArray[np.floating]) -> NDArray[np.floating]: # This is purely for testing and optimization sake # While calculating u^2, use u with einsum instead, as it is tad bit faster -def _skew_symmetrize_sq(vector: NDArray[np.floating]) -> NDArray[np.floating]: +def _skew_symmetrize_sq(vector: NDArray[np.float64]) -> NDArray[np.float64]: """ Generate the square of an orthogonal matrix from vector elements @@ -305,7 +305,7 @@ def _skew_symmetrize_sq(vector: NDArray[np.floating]) -> NDArray[np.floating]: # First generate array of [x^2, xy, xz, yx, y^2, yz, zx, zy, z^2] # across blocksize # This is slightly faster than doing v[np.newaxis,:,:] * v[:,np.newaxis,:] - products_xy: NDArray[np.floating] = np.einsum("ik,jk->ijk", vector, vector) + products_xy: NDArray[np.float64] = np.einsum("ik,jk->ijk", vector, vector) # No copy made here, as we do not change memory layout # products_xy = products_xy.reshape((dim * dim, -1)) @@ -338,8 +338,8 @@ def _skew_symmetrize_sq(vector: NDArray[np.floating]) -> NDArray[np.floating]: def _get_skew_symmetric_pair( - vector_collection: NDArray[np.floating], -) -> tuple[NDArray[np.floating], NDArray[np.floating]]: + vector_collection: NDArray[np.float64], +) -> tuple[NDArray[np.float64], NDArray[np.float64]]: """ Parameters @@ -355,7 +355,7 @@ def _get_skew_symmetric_pair( return u, u_sq -def _inv_skew_symmetrize(matrix: NDArray[np.floating]) -> NDArray[np.floating]: +def _inv_skew_symmetrize(matrix: NDArray[np.float64]) -> NDArray[np.float64]: """ Return the vector elements from a skew-symmetric matrix M diff --git a/elastica/_synchronize_periodic_boundary.py b/elastica/_synchronize_periodic_boundary.py index 49784f6c6..5af2b51e8 100644 --- a/elastica/_synchronize_periodic_boundary.py +++ b/elastica/_synchronize_periodic_boundary.py @@ -12,7 +12,7 @@ @njit(cache=True) # type: ignore def _synchronize_periodic_boundary_of_vector_collection( - input_array: NDArray[np.floating], periodic_idx: NDArray[np.floating] + input_array: NDArray[np.float64], periodic_idx: NDArray[np.float64] ) -> None: """ This function synchronizes the periodic boundaries of a vector collection. @@ -35,7 +35,7 @@ def _synchronize_periodic_boundary_of_vector_collection( @njit(cache=True) # type: ignore def _synchronize_periodic_boundary_of_matrix_collection( - input_array: NDArray[np.floating], periodic_idx: NDArray[np.floating] + input_array: NDArray[np.float64], periodic_idx: NDArray[np.float64] ) -> None: """ This function synchronizes the periodic boundaries of a matrix collection. @@ -61,7 +61,7 @@ def _synchronize_periodic_boundary_of_matrix_collection( @njit(cache=True) # type: ignore def _synchronize_periodic_boundary_of_scalar_collection( - input_array: NDArray[np.floating], periodic_idx: NDArray[np.floating] + input_array: NDArray[np.float64], periodic_idx: NDArray[np.float64] ) -> None: """ This function synchronizes the periodic boundaries of a scalar collection. @@ -92,7 +92,7 @@ class _ConstrainPeriodicBoundaries(ConstraintBase): def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) - def constrain_values(self, rod: RodType, time: np.floating) -> None: + def constrain_values(self, rod: RodType, time: np.float64) -> None: _synchronize_periodic_boundary_of_vector_collection( rod.position_collection, rod.periodic_boundary_nodes_idx ) @@ -100,7 +100,7 @@ def constrain_values(self, rod: RodType, time: np.floating) -> None: rod.director_collection, rod.periodic_boundary_elems_idx ) - def constrain_rates(self, rod: RodType, time: np.floating) -> None: + def constrain_rates(self, rod: RodType, time: np.float64) -> None: _synchronize_periodic_boundary_of_vector_collection( rod.velocity_collection, rod.periodic_boundary_nodes_idx ) diff --git a/elastica/boundary_conditions.py b/elastica/boundary_conditions.py index 9b5339937..e5ac86b4d 100644 --- a/elastica/boundary_conditions.py +++ b/elastica/boundary_conditions.py @@ -34,8 +34,8 @@ class ConstraintBase(ABC, Generic[S]): """ _system: S - _constrained_position_idx: NDArray[np.integer] - _constrained_director_idx: NDArray[np.integer] + _constrained_position_idx: NDArray[np.int32] + _constrained_director_idx: NDArray[np.int32] def __init__( self, @@ -74,7 +74,7 @@ def constrained_director_idx(self) -> np.ndarray: return self._constrained_director_idx @abstractmethod - def constrain_values(self, system: S, time: np.floating) -> None: + def constrain_values(self, system: S, time: np.float64) -> None: """ Constrain values (position and/or directors) of a rod object. @@ -88,7 +88,7 @@ def constrain_values(self, system: S, time: np.floating) -> None: pass @abstractmethod - def constrain_rates(self, system: S, time: np.floating) -> None: + def constrain_rates(self, system: S, time: np.float64) -> None: """ Constrain rates (velocity and/or omega) of a rod object. @@ -112,13 +112,13 @@ def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) def constrain_values( - self, system: "RodType | RigidBodyType", time: np.floating + self, system: "RodType | RigidBodyType", time: np.float64 ) -> None: """In FreeBC, this routine simply passes.""" pass def constrain_rates( - self, system: "RodType | RigidBodyType", time: np.floating + self, system: "RodType | RigidBodyType", time: np.float64 ) -> None: """In FreeBC, this routine simply passes.""" pass @@ -165,7 +165,7 @@ def __init__( self.fixed_directors_collection = np.array(fixed_directors) def constrain_values( - self, system: "RodType | RigidBodyType", time: np.floating + self, system: "RodType | RigidBodyType", time: np.float64 ) -> None: # system.position_collection[..., 0] = self.fixed_position # system.director_collection[..., 0] = self.fixed_directors @@ -177,7 +177,7 @@ def constrain_values( ) def constrain_rates( - self, system: "RodType | RigidBodyType", time: np.floating + self, system: "RodType | RigidBodyType", time: np.float64 ) -> None: # system.velocity_collection[..., 0] = 0.0 # system.omega_collection[..., 0] = 0.0 @@ -189,10 +189,10 @@ def constrain_rates( @staticmethod @njit(cache=True) # type: ignore def compute_constrain_values( - position_collection: NDArray[np.floating], - fixed_position_collection: NDArray[np.floating], - director_collection: NDArray[np.floating], - fixed_directors_collection: NDArray[np.floating], + position_collection: NDArray[np.float64], + fixed_position_collection: NDArray[np.float64], + director_collection: NDArray[np.float64], + fixed_directors_collection: NDArray[np.float64], ) -> None: """ Computes constrain values in numba njit decorator @@ -218,8 +218,8 @@ def compute_constrain_values( @staticmethod @njit(cache=True) # type: ignore def compute_constrain_rates( - velocity_collection: NDArray[np.floating], - omega_collection: NDArray[np.floating], + velocity_collection: NDArray[np.float64], + omega_collection: NDArray[np.float64], ) -> None: """ Compute contrain rates in numba njit decorator @@ -335,7 +335,7 @@ def __init__( self.rotational_constraint_selector = rotational_constraint_selector.astype(int) def constrain_values( - self, system: "RodType | RigidBodyType", time: np.floating + self, system: "RodType | RigidBodyType", time: np.float64 ) -> None: if self.constrained_position_idx.size: self.nb_constrain_translational_values( @@ -346,7 +346,7 @@ def constrain_values( ) def constrain_rates( - self, system: "RodType | RigidBodyType", time: np.floating + self, system: "RodType | RigidBodyType", time: np.float64 ) -> None: if self.constrained_position_idx.size: self.nb_constrain_translational_rates( @@ -365,10 +365,10 @@ def constrain_rates( @staticmethod @njit(cache=True) # type: ignore def nb_constrain_translational_values( - position_collection: NDArray[np.floating], - fixed_position_collection: NDArray[np.floating], - indices: NDArray[np.integer], - constraint_selector: NDArray[np.integer], + position_collection: NDArray[np.float64], + fixed_position_collection: NDArray[np.float64], + indices: NDArray[np.int32], + constraint_selector: NDArray[np.int32], ) -> None: """ Computes constrain values in numba njit decorator @@ -403,9 +403,9 @@ def nb_constrain_translational_values( @staticmethod @njit(cache=True) # type: ignore def nb_constrain_translational_rates( - velocity_collection: NDArray[np.floating], - indices: NDArray[np.integer], - constraint_selector: NDArray[np.integer], + velocity_collection: NDArray[np.float64], + indices: NDArray[np.int32], + constraint_selector: NDArray[np.int32], ) -> None: """ Compute constrain rates in numba njit decorator @@ -434,10 +434,10 @@ def nb_constrain_translational_rates( @staticmethod @njit(cache=True) # type: ignore def nb_constrain_rotational_rates( - director_collection: NDArray[np.floating], - omega_collection: NDArray[np.floating], - indices: NDArray[np.integer], - constraint_selector: NDArray[np.integer], + director_collection: NDArray[np.float64], + omega_collection: NDArray[np.float64], + indices: NDArray[np.int32], + constraint_selector: NDArray[np.int32], ) -> None: """ Compute constrain rates in numba njit decorator @@ -524,7 +524,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: ) def constrain_values( - self, system: "RodType | RigidBodyType", time: np.floating + self, system: "RodType | RigidBodyType", time: np.float64 ) -> None: if self.constrained_position_idx.size: self.nb_constrain_translational_values( @@ -540,7 +540,7 @@ def constrain_values( ) def constrain_rates( - self, system: "RodType | RigidBodyType", time: np.floating + self, system: "RodType | RigidBodyType", time: np.float64 ) -> None: if self.constrained_position_idx.size: self.nb_constrain_translational_rates( @@ -556,9 +556,9 @@ def constrain_rates( @staticmethod @njit(cache=True) # type: ignore def nb_constraint_rotational_values( - director_collection: NDArray[np.floating], - fixed_director_collection: NDArray[np.floating], - indices: NDArray[np.integer], + director_collection: NDArray[np.float64], + fixed_director_collection: NDArray[np.float64], + indices: NDArray[np.int32], ) -> None: """ Computes constrain values in numba njit decorator @@ -579,9 +579,9 @@ def nb_constraint_rotational_values( @staticmethod @njit(cache=True) # type: ignore def nb_constrain_translational_values( - position_collection: NDArray[np.floating], - fixed_position_collection: NDArray[np.floating], - indices: NDArray[np.integer], + position_collection: NDArray[np.float64], + fixed_position_collection: NDArray[np.float64], + indices: NDArray[np.int32], ) -> None: """ Computes constrain values in numba njit decorator @@ -602,7 +602,7 @@ def nb_constrain_translational_values( @staticmethod @njit(cache=True) # type: ignore def nb_constrain_translational_rates( - velocity_collection: NDArray[np.floating], indices: NDArray[np.integer] + velocity_collection: NDArray[np.float64], indices: NDArray[np.int32] ) -> None: """ Compute constrain rates in numba njit decorator @@ -624,7 +624,7 @@ def nb_constrain_translational_rates( @staticmethod @njit(cache=True) # type: ignore def nb_constrain_rotational_rates( - omega_collection: NDArray[np.floating], indices: NDArray[np.integer] + omega_collection: NDArray[np.float64], indices: NDArray[np.int32] ) -> None: """ Compute constrain rates in numba njit decorator @@ -681,10 +681,10 @@ class HelicalBucklingBC(ConstraintBase): def __init__( self, - position_start: NDArray[np.floating], - position_end: NDArray[np.floating], - director_start: NDArray[np.floating], - director_end: NDArray[np.floating], + position_start: NDArray[np.float64], + position_end: NDArray[np.float64], + director_start: NDArray[np.float64], + director_end: NDArray[np.float64], twisting_time: float, slack: float, number_of_rotations: float, @@ -746,7 +746,7 @@ def __init__( ) # rotation_matrix wants vectors 3,1 def constrain_values( - self, system: "RodType | RigidBodyType", time: np.floating + self, system: "RodType | RigidBodyType", time: np.float64 ) -> None: if time > self.twisting_time: system.position_collection[..., 0] = self.final_start_position @@ -756,7 +756,7 @@ def constrain_values( system.director_collection[..., -1] = self.final_end_directors def constrain_rates( - self, system: "RodType | RigidBodyType", time: np.floating + self, system: "RodType | RigidBodyType", time: np.float64 ) -> None: if time > self.twisting_time: system.velocity_collection[..., 0] = 0.0 diff --git a/elastica/callback_functions.py b/elastica/callback_functions.py index 8e89683de..0030b8b3b 100644 --- a/elastica/callback_functions.py +++ b/elastica/callback_functions.py @@ -32,7 +32,7 @@ def __init__(self) -> None: """ pass - def make_callback(self, system: T, time: np.floating, current_step: int) -> None: + def make_callback(self, system: T, time: np.float64, current_step: int) -> None: """ This method is called every time step. Users can define which parameters are called back and recorded. Also users @@ -81,7 +81,7 @@ def __init__(self, step_skip: int, callback_params: dict) -> None: self.callback_params = callback_params def make_callback( - self, system: "RodType | RigidBodyType", time: np.floating, current_step: int + self, system: "RodType | RigidBodyType", time: np.float64, current_step: int ) -> None: if current_step % self.sample_every == 0: @@ -176,7 +176,7 @@ def __init__( self.file_save_interval = file_save_interval # Data collector - self.buffer: dict[str, list[NDArray[np.floating] | np.floating | int]] = ( + self.buffer: dict[str, list[NDArray[np.float64] | np.float64 | int]] = ( defaultdict(list) ) self.buffer_size = 0 @@ -201,7 +201,7 @@ def __init__( self._ext = "pkl" def make_callback( - self, system: "RodType | RigidBodyType", time: np.floating, current_step: int + self, system: "RodType | RigidBodyType", time: np.float64, current_step: int ) -> None: """ diff --git a/elastica/collision/AABBCollection.py b/elastica/collision/AABBCollection.py index 9385dab40..c91b83b7a 100644 --- a/elastica/collision/AABBCollection.py +++ b/elastica/collision/AABBCollection.py @@ -11,8 +11,8 @@ class AABBCollection: def __init__( self, - elemental_position_collection: NDArray[np.floating], - dimension_collection: NDArray[np.floating], + elemental_position_collection: NDArray[np.float64], + dimension_collection: NDArray[np.float64], elements_per_aabb: int, ) -> None: """ @@ -85,8 +85,8 @@ def _update(self, aabb_collection: list["AABBCollection"]) -> None: def update( self, - elemental_position_collection: NDArray[np.floating], - dimension_collection: NDArray[np.floating], + elemental_position_collection: NDArray[np.float64], + dimension_collection: NDArray[np.float64], ) -> None: # Initialize the boxes for i in range(self.n_aabb): @@ -111,8 +111,8 @@ class AABBHierarchy: def __init__( self, - position_collection: NDArray[np.floating], - dimension_collection: NDArray[np.floating], + position_collection: NDArray[np.float64], + dimension_collection: NDArray[np.float64], avg_n_dofs_in_final_level: int, ) -> None: """ @@ -232,8 +232,8 @@ def n_aabbs_at_level(self, i: int) -> int: def update( self, - position_collection: NDArray[np.floating], - dimension_collection: NDArray[np.floating], + position_collection: NDArray[np.float64], + dimension_collection: NDArray[np.float64], ) -> None: # Update bottom level first, the first level entries n_aabbs_in_final_level = self.n_aabbs_at_level(self.n_levels - 1) @@ -278,7 +278,7 @@ def update( def are_aabb_intersecting( - first_aabb_collection: NDArray[np.floating], - second_aabb_collection: NDArray[np.floating], + first_aabb_collection: NDArray[np.float64], + second_aabb_collection: NDArray[np.float64], ) -> bool: return True diff --git a/elastica/contact_forces.py b/elastica/contact_forces.py index f2ee9ca78..70c7caf21 100644 --- a/elastica/contact_forces.py +++ b/elastica/contact_forces.py @@ -107,7 +107,7 @@ class RodRodContact(NoContact): """ - def __init__(self, k: np.floating, nu: np.floating) -> None: + def __init__(self, k: np.float64, nu: np.float64) -> None: """ Parameters ---------- @@ -538,8 +538,8 @@ def __init__( k: float, nu: float, slip_velocity_tol: float, - static_mu_array: NDArray[np.floating], - kinetic_mu_array: NDArray[np.floating], + static_mu_array: NDArray[np.float64], + kinetic_mu_array: NDArray[np.float64], ) -> None: """ Parameters diff --git a/elastica/contact_utils.py b/elastica/contact_utils.py index a67793a07..37887685b 100644 --- a/elastica/contact_utils.py +++ b/elastica/contact_utils.py @@ -11,7 +11,7 @@ @numba.njit(cache=True) # type: ignore -def _dot_product(a: Sequence[np.floating], b: Sequence[np.floating]) -> np.floating: +def _dot_product(a: Sequence[np.float64], b: Sequence[np.float64]) -> np.float64: total: np.float64 = np.float64(0.0) for i in range(3): total += a[i] * b[i] @@ -19,28 +19,28 @@ def _dot_product(a: Sequence[np.floating], b: Sequence[np.floating]) -> np.float @numba.njit(cache=True) # type: ignore -def _norm(a: Sequence[np.floating]) -> float: +def _norm(a: Sequence[np.float64]) -> float: return sqrt(_dot_product(a, a)) @numba.njit(cache=True) # type: ignore -def _clip(x: np.floating, low: np.floating, high: np.floating) -> np.floating: +def _clip(x: np.float64, low: np.float64, high: np.float64) -> np.float64: return max(low, min(x, high)) # Can this be made more efficient than 2 comp, 1 or? @numba.njit(cache=True) # type: ignore -def _out_of_bounds(x: np.floating, low: np.floating, high: np.floating) -> bool: - return (x < low) or (x > high) +def _out_of_bounds(x: np.float64, low: np.float64, high: np.float64) -> bool: + return bool((x < low) or (x > high)) @numba.njit(cache=True) # type: ignore def _find_min_dist( - x1: NDArray[np.floating], - e1: NDArray[np.floating], - x2: NDArray[np.floating], - e2: NDArray[np.floating], -) -> tuple[NDArray[np.floating], NDArray[np.floating], NDArray[np.floating]]: + x1: NDArray[np.float64], + e1: NDArray[np.float64], + x2: NDArray[np.float64], + e2: NDArray[np.float64], +) -> tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]]: e1e1 = _dot_product(e1, e1) e1e2 = _dot_product(e1, e2) e2e2 = _dot_product(e2, e2) @@ -107,7 +107,7 @@ def _find_min_dist( @numba.njit(cache=True) # type: ignore def _aabbs_not_intersecting( - aabb_one: NDArray[np.floating], aabb_two: NDArray[np.floating] + aabb_one: NDArray[np.float64], aabb_two: NDArray[np.float64] ) -> Literal[1, 0]: """Returns true if not intersecting else false""" if (aabb_one[0, 1] < aabb_two[0, 0]) | (aabb_one[0, 0] > aabb_two[0, 1]): @@ -122,13 +122,13 @@ def _aabbs_not_intersecting( @numba.njit(cache=True) # type: ignore def _prune_using_aabbs_rod_cylinder( - rod_one_position_collection: NDArray[np.floating], - rod_one_radius_collection: NDArray[np.floating], - rod_one_length_collection: NDArray[np.floating], - cylinder_position: NDArray[np.floating], - cylinder_director: NDArray[np.floating], - cylinder_radius: NDArray[np.floating], - cylinder_length: NDArray[np.floating], + rod_one_position_collection: NDArray[np.float64], + rod_one_radius_collection: NDArray[np.float64], + rod_one_length_collection: NDArray[np.float64], + cylinder_position: NDArray[np.float64], + cylinder_director: NDArray[np.float64], + cylinder_radius: NDArray[np.float64], + cylinder_length: NDArray[np.float64], ) -> Literal[1, 0]: max_possible_dimension = np.zeros((3,)) aabb_rod = np.empty((3, 2)) @@ -164,12 +164,12 @@ def _prune_using_aabbs_rod_cylinder( @numba.njit(cache=True) # type: ignore def _prune_using_aabbs_rod_rod( - rod_one_position_collection: NDArray[np.floating], - rod_one_radius_collection: NDArray[np.floating], - rod_one_length_collection: NDArray[np.floating], - rod_two_position_collection: NDArray[np.floating], - rod_two_radius_collection: NDArray[np.floating], - rod_two_length_collection: NDArray[np.floating], + rod_one_position_collection: NDArray[np.float64], + rod_one_radius_collection: NDArray[np.float64], + rod_one_length_collection: NDArray[np.float64], + rod_two_position_collection: NDArray[np.float64], + rod_two_radius_collection: NDArray[np.float64], + rod_two_length_collection: NDArray[np.float64], ) -> Literal[1, 0]: max_possible_dimension = np.zeros((3,)) aabb_rod_one = np.empty((3, 2)) @@ -202,12 +202,12 @@ def _prune_using_aabbs_rod_rod( @numba.njit(cache=True) # type: ignore def _prune_using_aabbs_rod_sphere( - rod_one_position_collection: NDArray[np.floating], - rod_one_radius_collection: NDArray[np.floating], - rod_one_length_collection: NDArray[np.floating], - sphere_position: NDArray[np.floating], - sphere_director: NDArray[np.floating], - sphere_radius: NDArray[np.floating], + rod_one_position_collection: NDArray[np.float64], + rod_one_radius_collection: NDArray[np.float64], + rod_one_length_collection: NDArray[np.float64], + sphere_position: NDArray[np.float64], + sphere_director: NDArray[np.float64], + sphere_radius: NDArray[np.float64], ) -> Literal[1, 0]: max_possible_dimension = np.zeros((3,)) aabb_rod = np.empty((3, 2)) @@ -241,8 +241,8 @@ def _prune_using_aabbs_rod_sphere( @numba.njit(cache=True) # type: ignore def _find_slipping_elements( - velocity_slip: NDArray[np.floating], velocity_threshold: np.floating -) -> NDArray[np.floating]: + velocity_slip: NDArray[np.float64], velocity_threshold: np.float64 +) -> NDArray[np.float64]: """ This function takes the velocity of elements and checks if they are larger than the threshold velocity. If the velocity of elements is larger than threshold velocity, that means those elements are slipping. @@ -283,7 +283,7 @@ def _find_slipping_elements( @numba.njit(cache=True) # type: ignore -def _node_to_element_mass_or_force(input: NDArray[np.floating]) -> NDArray[np.floating]: +def _node_to_element_mass_or_force(input: NDArray[np.float64]) -> NDArray[np.float64]: """ This function converts the mass/forces on rod nodes to elements, where special treatment is necessary at the ends. @@ -322,8 +322,8 @@ def _node_to_element_mass_or_force(input: NDArray[np.floating]) -> NDArray[np.fl @numba.njit(cache=True) # type: ignore def _elements_to_nodes_inplace( - vector_in_element_frame: NDArray[np.floating], - vector_in_node_frame: NDArray[np.floating], + vector_in_element_frame: NDArray[np.float64], + vector_in_node_frame: NDArray[np.float64], ) -> None: """ Updating nodal forces using the forces computed on elements @@ -348,8 +348,8 @@ def _elements_to_nodes_inplace( @numba.njit(cache=True) # type: ignore def _node_to_element_position( - node_position_collection: NDArray[np.floating], -) -> NDArray[np.floating]: + node_position_collection: NDArray[np.float64], +) -> NDArray[np.float64]: """ This function computes the position of the elements from the nodal values. @@ -396,8 +396,8 @@ def _node_to_element_position( @numba.njit(cache=True) # type: ignore def _node_to_element_velocity( - mass: NDArray[np.floating], node_velocity_collection: NDArray[np.floating] -) -> NDArray[np.floating]: + mass: NDArray[np.float64], node_velocity_collection: NDArray[np.float64] +) -> NDArray[np.float64]: """ This function computes the velocity of the elements from the nodal values. Uses the velocity of center of mass diff --git a/elastica/dissipation.py b/elastica/dissipation.py index f20c632e0..32ff48668 100644 --- a/elastica/dissipation.py +++ b/elastica/dissipation.py @@ -58,7 +58,7 @@ def system(self) -> T: return self._system @abstractmethod - def dampen_rates(self, system: T, time: np.floating) -> None: + def dampen_rates(self, system: T, time: np.float64) -> None: # TODO: In the future, we can remove rod and use self.system """ Dampen rates (velocity and/or omega) of a rod object. @@ -151,7 +151,7 @@ def __init__( * np.diagonal(self._system.inv_mass_second_moment_of_inertia).T ) - def dampen_rates(self, system: RodType, time: np.floating) -> None: + def dampen_rates(self, system: RodType, time: np.float64) -> None: system.velocity_collection[:] = ( system.velocity_collection * self.translational_damping_coefficient ) @@ -240,7 +240,7 @@ def __init__(self, filter_order: int, **kwargs: Any) -> None: self.omega_filter_term = np.zeros_like(self._system.omega_collection) self.filter_function = _filter_function_periodic_condition - def dampen_rates(self, system: RodType, time: np.floating) -> None: + def dampen_rates(self, system: RodType, time: np.float64) -> None: self.filter_function( system.velocity_collection, @@ -253,10 +253,10 @@ def dampen_rates(self, system: RodType, time: np.floating) -> None: @njit(cache=True) # type: ignore def _filter_function_periodic_condition_ring_rod( - velocity_collection: NDArray[np.floating], - velocity_filter_term: NDArray[np.floating], - omega_collection: NDArray[np.floating], - omega_filter_term: NDArray[np.floating], + velocity_collection: NDArray[np.float64], + velocity_filter_term: NDArray[np.float64], + omega_collection: NDArray[np.float64], + omega_filter_term: NDArray[np.float64], filter_order: int, ) -> None: blocksize = velocity_filter_term.shape[1] @@ -291,10 +291,10 @@ def _filter_function_periodic_condition_ring_rod( @njit(cache=True) # type: ignore def _filter_function_periodic_condition( - velocity_collection: NDArray[np.floating], - velocity_filter_term: NDArray[np.floating], - omega_collection: NDArray[np.floating], - omega_filter_term: NDArray[np.floating], + velocity_collection: NDArray[np.float64], + velocity_filter_term: NDArray[np.float64], + omega_collection: NDArray[np.float64], + omega_filter_term: NDArray[np.float64], filter_order: int, ) -> None: nb_filter_rate( @@ -311,8 +311,8 @@ def _filter_function_periodic_condition( @njit(cache=True) # type: ignore def nb_filter_rate( - rate_collection: NDArray[np.floating], - filter_term: NDArray[np.floating], + rate_collection: NDArray[np.float64], + filter_term: NDArray[np.float64], filter_order: int, ) -> None: """ diff --git a/elastica/external_forces.py b/elastica/external_forces.py index c3c49549e..6d377acb7 100644 --- a/elastica/external_forces.py +++ b/elastica/external_forces.py @@ -34,7 +34,7 @@ def __init__(self) -> None: """ pass - def apply_forces(self, system: S, time: np.floating = np.float64(0.0)) -> None: + def apply_forces(self, system: S, time: np.float64 = np.float64(0.0)) -> None: """Apply forces to a rod-like object. In NoForces class, this routine simply passes. @@ -49,7 +49,7 @@ def apply_forces(self, system: S, time: np.floating = np.float64(0.0)) -> None: """ pass - def apply_torques(self, system: S, time: np.floating = np.float64(0.0)) -> None: + def apply_torques(self, system: S, time: np.float64 = np.float64(0.0)) -> None: """Apply torques to a rod-like object. In NoForces class, this routine simply passes. @@ -77,7 +77,7 @@ class GravityForces(NoForces): """ def __init__( - self, acc_gravity: NDArray[np.floating] = np.array([0.0, -9.80665, 0.0]) + self, acc_gravity: NDArray[np.float64] = np.array([0.0, -9.80665, 0.0]) ) -> None: """ @@ -91,7 +91,7 @@ def __init__( self.acc_gravity = acc_gravity def apply_forces( - self, system: "RodType | RigidBodyType", time: np.floating = np.float64(0.0) + self, system: "RodType | RigidBodyType", time: np.float64 = np.float64(0.0) ) -> None: self.compute_gravity_forces( self.acc_gravity, system.mass, system.external_forces @@ -100,9 +100,9 @@ def apply_forces( @staticmethod @njit(cache=True) # type: ignore def compute_gravity_forces( - acc_gravity: NDArray[np.floating], - mass: NDArray[np.floating], - external_forces: NDArray[np.floating], + acc_gravity: NDArray[np.float64], + mass: NDArray[np.float64], + external_forces: NDArray[np.float64], ) -> None: """ This function add gravitational forces on the nodes. We are @@ -138,8 +138,8 @@ class EndpointForces(NoForces): def __init__( self, - start_force: NDArray[np.floating], - end_force: NDArray[np.floating], + start_force: NDArray[np.float64], + end_force: NDArray[np.float64], ramp_up_time: float, ) -> None: """ @@ -163,7 +163,7 @@ def __init__( self.ramp_up_time = np.float64(ramp_up_time) def apply_forces( - self, system: "RodType | RigidBodyType", time: np.floating = np.float64(0.0) + self, system: "RodType | RigidBodyType", time: np.float64 = np.float64(0.0) ) -> None: self.compute_end_point_forces( system.external_forces, @@ -176,11 +176,11 @@ def apply_forces( @staticmethod @njit(cache=True) # type: ignore def compute_end_point_forces( - external_forces: NDArray[np.floating], - start_force: NDArray[np.floating], - end_force: NDArray[np.floating], - time: np.floating, - ramp_up_time: np.floating, + external_forces: NDArray[np.float64], + start_force: NDArray[np.float64], + end_force: NDArray[np.float64], + time: np.float64, + ramp_up_time: np.float64, ) -> None: """ Compute end point forces that are applied on the rod using numba njit decorator. @@ -217,8 +217,8 @@ class UniformTorques(NoForces): def __init__( self, - torque: np.floating, - direction: NDArray[np.floating] = np.array([0.0, 0.0, 0.0]), + torque: np.float64, + direction: NDArray[np.float64] = np.array([0.0, 0.0, 0.0]), ) -> None: """ @@ -234,7 +234,7 @@ def __init__( self.torque = torque * direction def apply_torques( - self, system: "RodType | RigidBodyType", time: np.floating = np.float64(0.0) + self, system: "RodType | RigidBodyType", time: np.float64 = np.float64(0.0) ) -> None: n_elems = system.n_elems torque_on_one_element = ( @@ -257,8 +257,8 @@ class UniformForces(NoForces): def __init__( self, - force: np.floating, - direction: NDArray[np.floating] = np.array([0.0, 0.0, 0.0]), + force: np.float64, + direction: NDArray[np.float64] = np.array([0.0, 0.0, 0.0]), ) -> None: """ @@ -274,7 +274,7 @@ def __init__( self.force = (force * direction).reshape(3, 1) def apply_forces( - self, system: "RodType | RigidBodyType", time: np.floating = np.float64(0.0) + self, system: "RodType | RigidBodyType", time: np.float64 = np.float64(0.0) ) -> None: force_on_one_element = self.force / system.n_elems @@ -313,12 +313,12 @@ class MuscleTorques(NoForces): def __init__( self, base_length: float, # TODO: Is this necessary? - b_coeff: NDArray[np.floating], + b_coeff: NDArray[np.float64], period: float, wave_number: float, phase_shift: float, - direction: NDArray[np.floating], - rest_lengths: NDArray[np.floating], + direction: NDArray[np.float64], + rest_lengths: NDArray[np.float64], ramp_up_time: float, with_spline: bool = False, ) -> None: @@ -339,7 +339,7 @@ def __init__( Phase shift of traveling wave. direction: numpy.ndarray 1D (dim) array containing data with 'float' type. Muscle torque direction. - ramp_up_time: np.floating + ramp_up_time: np.float64 Applied muscle torques are ramped up until ramp up time. with_spline: boolean Option to use beta-spline. @@ -373,7 +373,7 @@ def __init__( self.my_spline = np.full_like(self.s, fill_value=1.0) def apply_torques( - self, system: "RodType | RigidBodyType", time: np.floating = np.float64(0.0) + self, system: "RodType | RigidBodyType", time: np.float64 = np.float64(0.0) ) -> None: self.compute_muscle_torques( time, @@ -392,15 +392,15 @@ def apply_torques( @njit(cache=True) # type: ignore def compute_muscle_torques( time: float, - my_spline: NDArray[np.floating], - s: np.floating, - angular_frequency: np.floating, - wave_number: np.floating, - phase_shift: np.floating, - ramp_up_time: np.floating, - direction: NDArray[np.floating], - director_collection: NDArray[np.floating], - external_torques: NDArray[np.floating], + my_spline: NDArray[np.float64], + s: np.float64, + angular_frequency: np.float64, + wave_number: np.float64, + phase_shift: np.float64, + ramp_up_time: np.float64, + direction: NDArray[np.float64], + director_collection: NDArray[np.float64], + external_torques: NDArray[np.float64], ) -> None: # Ramp up the muscle torque factor = np.float64(min(np.float64(1.0), time / ramp_up_time)) @@ -428,8 +428,8 @@ def compute_muscle_torques( @njit(cache=True) # type: ignore def inplace_addition( - external_force_or_torque: NDArray[np.floating], - force_or_torque: NDArray[np.floating], + external_force_or_torque: NDArray[np.float64], + force_or_torque: NDArray[np.float64], ) -> None: """ This function does inplace addition. First argument @@ -454,8 +454,8 @@ def inplace_addition( @njit(cache=True) # type: ignore def inplace_substraction( - external_force_or_torque: NDArray[np.floating], - force_or_torque: NDArray[np.floating], + external_force_or_torque: NDArray[np.float64], + force_or_torque: NDArray[np.float64], ) -> None: """ This function does inplace substraction. First argument @@ -508,8 +508,8 @@ def __init__( start_force_mag: float, end_force_mag: float, ramp_up_time: float = 0.0, - tangent_direction: NDArray[np.floating] = np.array([0, 0, 1]), - normal_direction: NDArray[np.floating] = np.array([0, 1, 0]), + tangent_direction: NDArray[np.float64] = np.array([0, 0, 1]), + normal_direction: NDArray[np.float64] = np.array([0, 1, 0]), ) -> None: """ @@ -541,7 +541,7 @@ def __init__( self.ramp_up_time = np.float64(ramp_up_time) def apply_forces( - self, system: "RodType | RigidBodyType", time: np.floating = np.float64(0.0) + self, system: "RodType | RigidBodyType", time: np.float64 = np.float64(0.0) ) -> None: if time < self.ramp_up_time: diff --git a/elastica/interaction.py b/elastica/interaction.py index 70ffcd34b..1a72650d0 100644 --- a/elastica/interaction.py +++ b/elastica/interaction.py @@ -48,8 +48,8 @@ def __init__( self, k: float, nu: float, - plane_origin: NDArray[np.floating], - plane_normal: NDArray[np.floating], + plane_origin: NDArray[np.float64], + plane_normal: NDArray[np.float64], ) -> None: """ @@ -72,9 +72,7 @@ def __init__( self.plane_origin = plane_origin.reshape(3, 1) self.plane_normal = plane_normal.reshape(3) - def apply_forces( - self, system: RodType, time: np.floating = np.float64(0.0) - ) -> None: + def apply_forces(self, system: RodType, time: np.float64 = np.float64(0.0)) -> None: """ In the case of contact with the plane, this function computes the plane reaction force on the element. @@ -146,11 +144,11 @@ def __init__( self, k: float, nu: float, - plane_origin: NDArray[np.floating], - plane_normal: NDArray[np.floating], + plane_origin: NDArray[np.float64], + plane_normal: NDArray[np.float64], slip_velocity_tol: float, - static_mu_array: NDArray[np.floating], - kinetic_mu_array: NDArray[np.floating], + static_mu_array: NDArray[np.float64], + kinetic_mu_array: NDArray[np.float64], ) -> None: """ @@ -191,7 +189,7 @@ def __init__( # kinetic and static friction should separate functions # for now putting them together to figure out common variables def apply_forces( - self, system: "RodType | RigidBodyType", time: np.floating = np.float64(0.0) + self, system: "RodType | RigidBodyType", time: np.float64 = np.float64(0.0) ) -> None: """ Call numba implementation to apply friction forces @@ -230,7 +228,7 @@ def apply_forces( # Slender body module @njit(cache=True) # type: ignore -def sum_over_elements(input: NDArray[np.floating]) -> np.floating: +def sum_over_elements(input: NDArray[np.float64]) -> np.float64: """ This function sums all elements of the input array. Using a Numba njit decorator shows better performance @@ -262,7 +260,7 @@ def sum_over_elements(input: NDArray[np.floating]) -> np.floating: This version: 513 ns ± 24.6 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each) """ - output: np.floating = np.float64(0.0) + output: np.float64 = np.float64(0.0) for i in range(input.shape[0]): output += input[i] @@ -271,13 +269,13 @@ def sum_over_elements(input: NDArray[np.floating]) -> np.floating: @njit(cache=True) # type: ignore def slender_body_forces( - tangents: NDArray[np.floating], - velocity_collection: NDArray[np.floating], - dynamic_viscosity: np.floating, - lengths: NDArray[np.floating], - radius: NDArray[np.floating], - mass: NDArray[np.floating], -) -> NDArray[np.floating]: + tangents: NDArray[np.float64], + velocity_collection: NDArray[np.float64], + dynamic_viscosity: np.float64, + lengths: NDArray[np.float64], + radius: NDArray[np.float64], + mass: NDArray[np.float64], +) -> NDArray[np.float64]: r""" This function computes hydrodynamic forces on a body using slender body theory. The below implementation is from Eq. 4.13 in Gazzola et al. RSoS. (2018). @@ -399,9 +397,7 @@ def __init__(self, dynamic_viscosity: float) -> None: super(SlenderBodyTheory, self).__init__() self.dynamic_viscosity = np.float64(dynamic_viscosity) - def apply_forces( - self, system: RodType, time: np.floating = np.float64(0.0) - ) -> None: + def apply_forces(self, system: RodType, time: np.float64 = np.float64(0.0)) -> None: """ This function applies hydrodynamic forces on body using the slender body theory given in @@ -431,8 +427,8 @@ def __init__( self, k: float, nu: float, - plane_origin: NDArray[np.floating], - plane_normal: NDArray[np.floating], + plane_origin: NDArray[np.float64], + plane_normal: NDArray[np.float64], ) -> None: self.k = np.float64(k) self.nu = np.float64(nu) @@ -441,7 +437,7 @@ def __init__( self.plane_normal = plane_normal.reshape(3) def apply_forces( - self, system: RigidBodyType, time: np.floating = np.float64(0.0) + self, system: RigidBodyType, time: np.float64 = np.float64(0.0) ) -> None: """ This function computes the plane force response on the rigid body, in the diff --git a/elastica/joint.py b/elastica/joint.py index 2788b5019..1cf2dd3eb 100644 --- a/elastica/joint.py +++ b/elastica/joint.py @@ -143,7 +143,7 @@ def __init__( k: float, nu: float, kt: float, - normal_direction: NDArray[np.floating], + normal_direction: NDArray[np.float64], ) -> None: """ @@ -238,7 +238,7 @@ def __init__( nu: float, kt: float, nut: float = 0.0, - rest_rotation_matrix: NDArray[np.floating] | None = None, + rest_rotation_matrix: NDArray[np.float64] | None = None, ) -> None: """ @@ -338,7 +338,7 @@ def get_relative_rotation_two_systems( index_one: ConnectionIndex, system_two: "RodType | RigidBodyType", index_two: ConnectionIndex, -) -> NDArray[np.floating]: +) -> NDArray[np.float64]: """ Compute the relative rotation matrix C_12 between system one and system two at the specified elements. diff --git a/elastica/memory_block/memory_block_rigid_body.py b/elastica/memory_block/memory_block_rigid_body.py index 38ebdb8f1..12c1caeab 100644 --- a/elastica/memory_block/memory_block_rigid_body.py +++ b/elastica/memory_block/memory_block_rigid_body.py @@ -15,7 +15,7 @@ def __init__( self.n_systems = len(systems) self.n_elems = self.n_systems self.n_nodes = self.n_elems - self.system_idx_list = np.array(system_idx_list, dtype=np.int64) + self.system_idx_list = np.array(system_idx_list, dtype=np.int32) # Allocate block structure using system collection. self._allocate_block_variables_scalars(systems) diff --git a/elastica/memory_block/memory_block_rod.py b/elastica/memory_block/memory_block_rod.py index 6cc07a086..a363f3d3c 100644 --- a/elastica/memory_block/memory_block_rod.py +++ b/elastica/memory_block/memory_block_rod.py @@ -53,20 +53,20 @@ def __init__( # Sorted systems systems = system_straight_rod + system_ring_rod self.system_idx_list = np.array( - system_idx_list_straight_rod + system_idx_list_ring_rod, dtype=np.int64 + system_idx_list_straight_rod + system_idx_list_ring_rod, dtype=np.int32 ) n_elems_straight_rods = np.array( - [x.n_elems for x in system_straight_rod], dtype=np.int64 + [x.n_elems for x in system_straight_rod], dtype=np.int32 ) n_elems_ring_rods = np.array( - [x.n_elems for x in system_ring_rod], dtype=np.int64 + [x.n_elems for x in system_ring_rod], dtype=np.int32 ) n_straight_rods: int = len(system_straight_rod) n_ring_rods: int = len(system_ring_rod) - # self.n_elems_in_rods = np.array([x.n_elems for x in systems], dtype=np.int) + # self.n_elems_in_rods = np.array([x.n_elems for x in systems], dtype=np.int32) self.n_elems_in_rods = np.hstack((n_elems_straight_rods, n_elems_ring_rods + 2)) self.n_rods = len(systems) ( diff --git a/elastica/memory_block/utils.py b/elastica/memory_block/utils.py index c7219e405..47c404898 100644 --- a/elastica/memory_block/utils.py +++ b/elastica/memory_block/utils.py @@ -1,15 +1,16 @@ __doc__ = """Create block-structure class for collection of Cosserat rod systems.""" import numpy as np import numpy.typing as npt +from numpy.typing import NDArray def make_block_memory_metadata( - n_elems_in_rods: npt.NDArray[np.integer], + n_elems_in_rods: NDArray[np.int32], ) -> tuple[ int, - npt.NDArray[np.integer], - npt.NDArray[np.integer], - npt.NDArray[np.integer], + NDArray[np.int32], + NDArray[np.int32], + NDArray[np.int32], ]: """ This function, takes number of elements of each rod as a numpy array and computes, @@ -17,22 +18,21 @@ def make_block_memory_metadata( Parameters ---------- - n_elems_in_rods: npt.NDArray + n_elems_in_rods: NDArray An integer array containing the number of elements in each of the n rod. Returns ------- - n_elems_with_ghosts: int64 + n_elems_with_ghosts: np.int32 Total number of elements with ghost elements included. There are two ghost elements between each pair of two rods adjacent in memory block. - ghost_nodes_idx: ndarray + ghost_nodes_idx: NDArray[np.int32] An integer array of length n - 1 containing the indices of ghost nodes in memory block. - ghost_elements_idx: npt.NDArray + ghost_elements_idx: NDArray[np.int32] An integer array of length 2 * (n - 1) containing the indices of ghost elements in memory block. - ghost_voronoi_idx: npt.NDArray + ghost_voronoi_idx: NDArray[np.int32] An integer array of length 2 * (n - 1) containing the indices of ghost Voronoi nodes in memory block. """ - n_nodes_in_rods = n_elems_in_rods + 1 n_rods = n_elems_in_rods.shape[0] @@ -43,29 +43,29 @@ def make_block_memory_metadata( # Gap between two rods have three ghost voronois : comes out to n_nodes_with_ghosts - 2 # n_voronoi_with_ghosts = np.sum(n_voronois_in_rods) + 3 * (n_rods - 1) - ghost_nodes_idx = np.cumsum(n_nodes_in_rods[:-1], dtype=np.int64) + ghost_nodes_idx = np.cumsum(n_nodes_in_rods[:-1], dtype=np.int32) # Add [0, 1, 2, ... n_rods-2] to the ghost_nodes idx to accommodate miscounting - ghost_nodes_idx += np.arange(0, n_rods - 1, dtype=np.int64) + ghost_nodes_idx += np.arange(n_rods - 1) - ghost_elems_idx = np.zeros((2 * (n_rods - 1),), dtype=np.int64) + ghost_elems_idx = np.zeros((2 * (n_rods - 1),), dtype=np.int32) ghost_elems_idx[::2] = ghost_nodes_idx - 1 ghost_elems_idx[1::2] = ghost_nodes_idx.copy() - ghost_voronoi_idx = np.zeros((3 * (n_rods - 1),), dtype=np.int64) + ghost_voronoi_idx = np.zeros((3 * (n_rods - 1),), dtype=np.int32) ghost_voronoi_idx[::3] = ghost_nodes_idx - 2 ghost_voronoi_idx[1::3] = ghost_nodes_idx - 1 ghost_voronoi_idx[2::3] = ghost_nodes_idx.copy() - return n_elems_with_ghosts, ghost_nodes_idx, ghost_elems_idx, ghost_voronoi_idx + return int(n_elems_with_ghosts), ghost_nodes_idx, ghost_elems_idx, ghost_voronoi_idx def make_block_memory_periodic_boundary_metadata( - n_elems_in_rods: npt.NDArray[np.integer], + n_elems_in_rods: NDArray[np.int32], ) -> tuple[ - npt.NDArray[np.integer], - npt.NDArray[np.integer], - npt.NDArray[np.integer], - npt.NDArray[np.integer], + NDArray[np.int32], + NDArray[np.int32], + NDArray[np.int32], + NDArray[np.int32], ]: """ This function, takes the number of elements of ring rods and computes the periodic boundary node, @@ -73,7 +73,7 @@ def make_block_memory_periodic_boundary_metadata( Parameters ---------- - n_elems_in_rods : npt.NDArray + n_elems_in_rods : NDArray 1D (n_ring_rods,) array containing data with 'int' type. Elements of this array contains total number of elements of one rod, including periodic boundary elements. @@ -81,24 +81,24 @@ def make_block_memory_periodic_boundary_metadata( ------- n_elems - periodic_boundary_node : npt.NDArray + periodic_boundary_node : NDArray 2D (2, n_periodic_boundary_nodes) array containing data with 'int' type. Vector containing periodic boundary elements index. First dimension is the periodic boundary index, second dimension is the referenced cell index. - periodic_boundary_elems_idx : npt.NDArray + periodic_boundary_elems_idx : NDArray 2D (2, n_periodic_boundary_elems) array containing data with 'int' type. Vector containing periodic boundary nodes index. First dimension is the periodic boundary index, second dimension is the referenced cell index. - periodic_boundary_voronoi_idx : npt.NDArray + periodic_boundary_voronoi_idx : NDArray 2D (2, n_periodic_boundary_voronoi) array containing data with 'int' type. Vector containing periodic boundary voronoi index. First dimension is the periodic boundary index, second dimension is the referenced cell index. """ - n_elem: npt.NDArray[np.integer] = n_elems_in_rods.copy() + n_elem: NDArray[np.int32] = n_elems_in_rods.copy() n_rods = n_elems_in_rods.shape[0] - periodic_boundary_node_idx = np.zeros((2, 3 * n_rods), dtype=np.int64) + periodic_boundary_node_idx = np.zeros((2, 3 * n_rods), dtype=np.int32) # count ghost nodes, first rod does not have a ghost node at the start, so exclude first rod. periodic_boundary_node_idx[0, 0::3][1:] = 1 # This is for the first periodic node at the end @@ -107,9 +107,7 @@ def make_block_memory_periodic_boundary_metadata( periodic_boundary_node_idx[0, 2::3] = 1 periodic_boundary_node_idx[0, :] = np.cumsum(periodic_boundary_node_idx[0, :]) # Add [0, 1, 2, ..., n_rods] to the periodic boundary nodes to accommodate miscounting - periodic_boundary_node_idx[0, :] += np.repeat( - np.arange(0, n_rods, dtype=np.int64), 3 - ) + periodic_boundary_node_idx[0, :] += np.repeat(np.arange(n_rods), 3) # Now fill the reference node idx, to copy and correct periodic boundary nodes # First fill with the reference node idx of the first periodic node. This is the last node of the actual rod # (without ghost and periodic nodes). @@ -121,16 +119,14 @@ def make_block_memory_periodic_boundary_metadata( # (without ghost and periodic nodes). periodic_boundary_node_idx[1, 2::3] = periodic_boundary_node_idx[0, 0::3] + 2 - periodic_boundary_elems_idx = np.zeros((2, 2 * n_rods), dtype=np.int64) + periodic_boundary_elems_idx = np.zeros((2, 2 * n_rods), dtype=np.int32) # count ghost elems, first rod does not have a ghost elem at the start, so exclude first rod. periodic_boundary_elems_idx[0, 0::2][1:] = 2 # This is for the first periodic elem at the end periodic_boundary_elems_idx[0, 1::2] = 1 + n_elem periodic_boundary_elems_idx[0, :] = np.cumsum(periodic_boundary_elems_idx[0, :]) # Add [0, 1, 2, ..., n_rods] to the periodic boundary elems to accommodate miscounting - periodic_boundary_elems_idx[0, :] += np.repeat( - np.arange(0, n_rods, dtype=np.int64), 2 - ) + periodic_boundary_elems_idx[0, :] += np.repeat(np.arange(n_rods), 2) # Now fill the reference element idx, to copy and correct periodic boundary elements # First fill with the reference element idx of the first periodic element. This is the last element of the actual # rod @@ -141,16 +137,14 @@ def make_block_memory_periodic_boundary_metadata( # (without ghost and periodic elements). periodic_boundary_elems_idx[1, 1::2] = periodic_boundary_elems_idx[0, 0::2] + 1 - periodic_boundary_voronoi_idx = np.zeros((2, n_rods), dtype=np.int64) + periodic_boundary_voronoi_idx = np.zeros((2, n_rods), dtype=np.int32) # count ghost voronoi, first rod does not have a ghost voronoi at the start, so exclude first rod. periodic_boundary_voronoi_idx[0, 0::1][1:] = 3 # This is for the first periodic voronoi at the end periodic_boundary_voronoi_idx[0, 1:] += n_elem[:-1] periodic_boundary_voronoi_idx[0, :] = np.cumsum(periodic_boundary_voronoi_idx[0, :]) # Add [0, 1, 2, ..., n_rods] to the periodic boundary voronoi to accommodate miscounting - periodic_boundary_voronoi_idx[0, :] += np.repeat( - np.arange(0, n_rods, dtype=np.int64), 1 - ) + periodic_boundary_voronoi_idx[0, :] += np.repeat(np.arange(0, n_rods), 1) # Now fill the reference voronoi idx, to copy and correct periodic boundary voronoi # Fill with the reference voronoi idx of the periodic voronoi. This is the last voronoi of the actual rod # (without ghost and periodic voronoi). diff --git a/elastica/mesh/protocol.py b/elastica/mesh/protocol.py index 766651fd5..f4dd2fd34 100644 --- a/elastica/mesh/protocol.py +++ b/elastica/mesh/protocol.py @@ -5,6 +5,6 @@ class MeshProtocol(Protocol): - faces: NDArray[np.floating] - face_centers: NDArray[np.floating] - face_normals: NDArray[np.floating] + faces: NDArray[np.float64] + face_centers: NDArray[np.float64] + face_normals: NDArray[np.float64] diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index faf922a1a..10c111ce6 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -203,25 +203,25 @@ def finalize(self) -> None: del self._feature_group_finalize @final - def synchronize(self, time: np.floating) -> None: + def synchronize(self, time: np.float64) -> None: # Collection call _feature_group_synchronize for func in self._feature_group_synchronize: func(time=time) @final - def constrain_values(self, time: np.floating) -> None: + def constrain_values(self, time: np.float64) -> None: # Collection call _feature_group_constrain_values for func in self._feature_group_constrain_values: func(time=time) @final - def constrain_rates(self, time: np.floating) -> None: + def constrain_rates(self, time: np.float64) -> None: # Collection call _feature_group_constrain_rates for func in self._feature_group_constrain_rates: func(time=time) @final - def apply_callbacks(self, time: np.floating, current_step: int) -> None: + def apply_callbacks(self, time: np.float64, current_step: int) -> None: # Collection call _feature_group_callback for func in self._feature_group_callback: func(time=time, current_step=current_step) diff --git a/elastica/modules/callbacks.py b/elastica/modules/callbacks.py index bf67357f3..4a299da0a 100644 --- a/elastica/modules/callbacks.py +++ b/elastica/modules/callbacks.py @@ -73,7 +73,7 @@ def _finalize_callback(self: SystemCollectionProtocol) -> None: def _callback_execution( self: SystemCollectionProtocol, - time: np.floating, + time: np.float64, current_step: int, ) -> None: for sys_id, callback in self._callback_operators: diff --git a/elastica/modules/connections.py b/elastica/modules/connections.py index 71b15834d..c55d5beaf 100644 --- a/elastica/modules/connections.py +++ b/elastica/modules/connections.py @@ -88,7 +88,7 @@ def _finalize_connections(self: SystemCollectionProtocol) -> None: # to apply the connections to. def apply_forces_and_torques( - time: np.floating, + time: np.float64, connect_instance: FreeJoint, system_one: "RodType | RigidBodyType", first_connect_idx: ConnectionIndex, diff --git a/elastica/modules/constraints.py b/elastica/modules/constraints.py index 75ec8e4dd..4a2c58407 100644 --- a/elastica/modules/constraints.py +++ b/elastica/modules/constraints.py @@ -107,11 +107,11 @@ def _finalize_constraints(self: SystemCollectionProtocol) -> None: self._constrain_values(time=np.float64(0.0)) self._constrain_rates(time=np.float64(0.0)) - def _constrain_values(self: SystemCollectionProtocol, time: np.floating) -> None: + def _constrain_values(self: SystemCollectionProtocol, time: np.float64) -> None: for sys_id, constraint in self._constraints_operators: constraint.constrain_values(self._systems[sys_id], time) - def _constrain_rates(self: SystemCollectionProtocol, time: np.floating) -> None: + def _constrain_rates(self: SystemCollectionProtocol, time: np.float64) -> None: for sys_id, constraint in self._constraints_operators: constraint.constrain_rates(self._systems[sys_id], time) diff --git a/elastica/modules/contact.py b/elastica/modules/contact.py index f52bdc31e..d30136355 100644 --- a/elastica/modules/contact.py +++ b/elastica/modules/contact.py @@ -80,7 +80,7 @@ def _finalize_contact(self: SystemCollectionProtocol) -> None: # to apply the contacts to def apply_contact( - time: np.floating, + time: np.float64, contact_instance: NoContact, first_sys_idx: SystemIdxType, second_sys_idx: SystemIdxType, diff --git a/elastica/modules/damping.py b/elastica/modules/damping.py index a3a2615e5..646e1a21d 100644 --- a/elastica/modules/damping.py +++ b/elastica/modules/damping.py @@ -76,7 +76,7 @@ def _finalize_dampers(self: SystemCollectionProtocol) -> None: # to sort dampers. self._damping_operators.sort(key=lambda x: x[0]) - def _dampen_rates(self: SystemCollectionProtocol, time: np.floating) -> None: + def _dampen_rates(self: SystemCollectionProtocol, time: np.float64) -> None: for sys_id, damper in self._damping_operators: damper.dampen_rates(self._systems[sys_id], time) diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index 5de97c7e2..b9a52ec7c 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -47,22 +47,22 @@ def __getitem__(self, i: slice | int) -> "list[SystemType] | SystemType": ... @property def _feature_group_synchronize(self) -> OperatorGroupFIFO: ... - def synchronize(self, time: np.floating) -> None: ... + def synchronize(self, time: np.float64) -> None: ... @property def _feature_group_constrain_values(self) -> list[OperatorType]: ... - def constrain_values(self, time: np.floating) -> None: ... + def constrain_values(self, time: np.float64) -> None: ... @property def _feature_group_constrain_rates(self) -> list[OperatorType]: ... - def constrain_rates(self, time: np.floating) -> None: ... + def constrain_rates(self, time: np.float64) -> None: ... @property def _feature_group_callback(self) -> list[OperatorCallbackType]: ... - def apply_callbacks(self, time: np.floating, current_step: int) -> None: ... + def apply_callbacks(self, time: np.float64, current_step: int) -> None: ... @property def _feature_group_finalize(self) -> list[OperatorFinalizeType]: ... @@ -98,7 +98,7 @@ def collect_diagnostics(self, system: SystemType) -> ModuleProtocol: @abstractmethod def _callback_execution( - self, time: np.floating, current_step: int, *args: Any, **kwargs: Any + self, time: np.float64, current_step: int, *args: Any, **kwargs: Any ) -> None: raise NotImplementedError @@ -112,11 +112,11 @@ def constrain(self, system: SystemType) -> ModuleProtocol: raise NotImplementedError @abstractmethod - def _constrain_values(self, time: np.floating) -> None: + def _constrain_values(self, time: np.float64) -> None: raise NotImplementedError @abstractmethod - def _constrain_rates(self, time: np.floating) -> None: + def _constrain_rates(self, time: np.float64) -> None: raise NotImplementedError # Forcing API @@ -147,5 +147,5 @@ def dampen(self, system: SystemType) -> ModuleProtocol: raise NotImplementedError @abstractmethod - def _dampen_rates(self, time: np.floating) -> None: + def _dampen_rates(self, time: np.float64) -> None: raise NotImplementedError diff --git a/elastica/reset_functions_for_block_structure/_reset_ghost_vector_or_scalar.py b/elastica/reset_functions_for_block_structure/_reset_ghost_vector_or_scalar.py index 18e5d129c..1759a5bac 100644 --- a/elastica/reset_functions_for_block_structure/_reset_ghost_vector_or_scalar.py +++ b/elastica/reset_functions_for_block_structure/_reset_ghost_vector_or_scalar.py @@ -7,18 +7,18 @@ @njit(cache=True) # type: ignore def _reset_vector_ghost( - input: NDArray[np.floating], - ghost_idx: NDArray[np.integer], - reset_value: np.floating = np.float64(0.0), + input: NDArray[np.float64], + ghost_idx: NDArray[np.int32], + reset_value: np.float64 = np.float64(0.0), ): """ This function resets the ghost of an input vector collection. Default reset value is 0.0. Parameters ---------- - input : NDArray[np.floating] - ghost_idx : NDArray[np.integer] - reset_value : np.floating + input : NDArray[np.float64] + ghost_idx : NDArray[np.int32] + reset_value : np.float64 Returns ------- @@ -47,18 +47,18 @@ def _reset_vector_ghost( @njit(cache=True) # type: ignore def _reset_scalar_ghost( - input: NDArray[np.floating], - ghost_idx: NDArray[np.integer], - reset_value: np.floating = np.float64(0.0), + input: NDArray[np.float64], + ghost_idx: NDArray[np.int32], + reset_value: np.float64 = np.float64(0.0), ): """ This function resets the ghost of a scalar collection. Default reset value is 0.0. Parameters ---------- - input : NDArray[np.floating] - ghost_idx : NDArray[np.integer] - reset_value : np.floating + input : NDArray[np.float64] + ghost_idx : NDArray[np.int32] + reset_value : np.float64 Returns ------- diff --git a/elastica/restart.py b/elastica/restart.py index cd31d520f..d6f9475b2 100644 --- a/elastica/restart.py +++ b/elastica/restart.py @@ -30,7 +30,7 @@ def all_equal(iterable: Iterable[Any]) -> bool: def save_state( simulator: Iterable, directory: str = "", - time: np.floating = np.float64(0.0), + time: np.float64 = np.float64(0.0), verbose: bool = False, ) -> None: """ diff --git a/elastica/rigidbody/cylinder.py b/elastica/rigidbody/cylinder.py index 6dab8384c..f1e8e83d5 100644 --- a/elastica/rigidbody/cylinder.py +++ b/elastica/rigidbody/cylinder.py @@ -14,9 +14,9 @@ class Cylinder(RigidBodyBase): def __init__( self, - start: NDArray[np.floating], - direction: NDArray[np.floating], - normal: NDArray[np.floating], + start: NDArray[np.float64], + direction: NDArray[np.float64], + normal: NDArray[np.float64], base_length: float, base_radius: float, density: float, @@ -26,9 +26,9 @@ def __init__( Parameters ---------- - start : NDArray[np.floating] - direction : NDArray[np.floating] - normal : NDArray[np.floating] + start : NDArray[np.float64] + direction : NDArray[np.float64] + normal : NDArray[np.float64] base_length : float base_radius : float density : float @@ -36,7 +36,7 @@ def __init__( # FIXME: Refactor def assert_check_array_size( - to_check: NDArray[np.floating], name: str, expected: int = 3 + to_check: NDArray[np.float64], name: str, expected: int = 3 ) -> None: array_size = to_check.size assert array_size == expected, ( diff --git a/elastica/rigidbody/mesh_rigid_body.py b/elastica/rigidbody/mesh_rigid_body.py index f72925b01..e5cf3f50d 100644 --- a/elastica/rigidbody/mesh_rigid_body.py +++ b/elastica/rigidbody/mesh_rigid_body.py @@ -14,10 +14,10 @@ class MeshRigidBody(RigidBodyBase): def __init__( self, mesh: MeshType, - center_of_mass: NDArray[np.floating], - mass_second_moment_of_inertia: NDArray[np.floating], - density: np.floating, - volume: np.floating, + center_of_mass: NDArray[np.float64], + mass_second_moment_of_inertia: NDArray[np.float64], + density: np.float64, + volume: np.float64, ) -> None: """ Mesh rigid body initializer. diff --git a/elastica/rigidbody/protocol.py b/elastica/rigidbody/protocol.py index 22bc049d9..7e6f0664e 100644 --- a/elastica/rigidbody/protocol.py +++ b/elastica/rigidbody/protocol.py @@ -8,11 +8,11 @@ class RigidBodyProtocol(SystemProtocol, SlenderBodyGeometryProtocol, Protocol): - mass: np.floating - volume: np.floating - length: np.floating - tangents: NDArray[np.floating] - radius: np.floating + mass: np.float64 + volume: np.float64 + length: np.float64 + tangents: NDArray[np.float64] + radius: np.float64 - mass_second_moment_of_inertia: NDArray[np.floating] - inv_mass_second_moment_of_inertia: NDArray[np.floating] + mass_second_moment_of_inertia: NDArray[np.float64] + inv_mass_second_moment_of_inertia: NDArray[np.float64] diff --git a/elastica/rigidbody/rigid_body.py b/elastica/rigidbody/rigid_body.py index a7252e049..999208546 100644 --- a/elastica/rigidbody/rigid_body.py +++ b/elastica/rigidbody/rigid_body.py @@ -27,29 +27,29 @@ def __init__(self) -> None: self.n_elems: int = 1 self.n_nodes: int = 1 - self.position_collection: NDArray[np.floating] - self.velocity_collection: NDArray[np.floating] - self.acceleration_collection: NDArray[np.floating] - self.omega_collection: NDArray[np.floating] - self.alpha_collection: NDArray[np.floating] - self.director_collection: NDArray[np.floating] + self.position_collection: NDArray[np.float64] + self.velocity_collection: NDArray[np.float64] + self.acceleration_collection: NDArray[np.float64] + self.omega_collection: NDArray[np.float64] + self.alpha_collection: NDArray[np.float64] + self.director_collection: NDArray[np.float64] - self.external_forces: NDArray[np.floating] - self.external_torques: NDArray[np.floating] + self.external_forces: NDArray[np.float64] + self.external_torques: NDArray[np.float64] - self.internal_forces: NDArray[np.floating] - self.internal_torques: NDArray[np.floating] + self.internal_forces: NDArray[np.float64] + self.internal_torques: NDArray[np.float64] - self.mass: np.floating - self.volume: np.floating - self.length: np.floating - self.tangents: NDArray[np.floating] - self.radius: np.floating + self.mass: np.float64 + self.volume: np.float64 + self.length: np.float64 + self.tangents: NDArray[np.float64] + self.radius: np.float64 - self.mass_second_moment_of_inertia: NDArray[np.floating] - self.inv_mass_second_moment_of_inertia: NDArray[np.floating] + self.mass_second_moment_of_inertia: NDArray[np.float64] + self.inv_mass_second_moment_of_inertia: NDArray[np.float64] - def update_accelerations(self, time: np.floating) -> None: + def update_accelerations(self, time: np.float64) -> None: np.copyto( self.acceleration_collection, (self.external_forces) / self.mass, @@ -71,24 +71,24 @@ def update_accelerations(self, time: np.floating) -> None: ), ) - def zeroed_out_external_forces_and_torques(self, time: np.floating) -> None: + def zeroed_out_external_forces_and_torques(self, time: np.float64) -> None: # Reset forces and torques self.external_forces *= 0.0 self.external_torques *= 0.0 - def compute_internal_forces_and_torques(self, time: np.floating) -> None: + def compute_internal_forces_and_torques(self, time: np.float64) -> None: """ For rigid body, there is no internal forces and torques """ pass - def compute_position_center_of_mass(self) -> NDArray[np.floating]: + def compute_position_center_of_mass(self) -> NDArray[np.float64]: """ Return positional center of mass """ return self.position_collection[..., 0].copy() - def compute_translational_energy(self) -> NDArray[np.floating]: + def compute_translational_energy(self) -> NDArray[np.float64]: """ Return translational energy """ @@ -100,7 +100,7 @@ def compute_translational_energy(self) -> NDArray[np.floating]: ) ) - def compute_rotational_energy(self) -> NDArray[np.floating]: + def compute_rotational_energy(self) -> NDArray[np.float64]: """ Return rotational energy """ diff --git a/elastica/rigidbody/sphere.py b/elastica/rigidbody/sphere.py index 72c4481b8..5f87a42a0 100644 --- a/elastica/rigidbody/sphere.py +++ b/elastica/rigidbody/sphere.py @@ -14,7 +14,7 @@ class Sphere(RigidBodyBase): def __init__( self, - center: NDArray[np.floating], + center: NDArray[np.float64], base_radius: float, density: float, ) -> None: @@ -23,7 +23,7 @@ def __init__( Parameters ---------- - center : NDArray[np.floating] + center : NDArray[np.float64] base_radius : float density : float """ diff --git a/elastica/rod/cosserat_rod.py b/elastica/rod/cosserat_rod.py index ea79854a1..8ac1d081b 100644 --- a/elastica/rod/cosserat_rod.py +++ b/elastica/rod/cosserat_rod.py @@ -32,7 +32,7 @@ @functools.lru_cache(maxsize=1) -def _get_z_vector() -> NDArray[np.floating]: +def _get_z_vector() -> NDArray[np.float64]: return np.array([0.0, 0.0, 1.0]).reshape(3, -1) @@ -79,73 +79,73 @@ class CosseratRod(RodBase, KnotTheory): ---------- n_elems: int The number of elements of the rod. - position_collection: NDArray[np.floating] + position_collection: NDArray[np.float64] 2D (dim, n_nodes) array containing data with 'float' type. Array containing node position vectors. - velocity_collection: NDArray[np.floating] + velocity_collection: NDArray[np.float64] 2D (dim, n_nodes) array containing data with 'float' type. Array containing node velocity vectors. - acceleration_collection: NDArray[np.floating] + acceleration_collection: NDArray[np.float64] 2D (dim, n_nodes) array containing data with 'float' type. Array containing node acceleration vectors. - omega_collection: NDArray[np.floating] + omega_collection: NDArray[np.float64] 2D (dim, n_elems) array containing data with 'float' type. Array containing element angular velocity vectors. - alpha_collection: NDArray[np.floating] + alpha_collection: NDArray[np.float64] 2D (dim, n_elems) array containing data with 'float' type. Array contining element angular acceleration vectors. - director_collection: NDArray[np.floating] + director_collection: NDArray[np.float64] 3D (dim, dim, n_elems) array containing data with 'float' type. Array containing element director matrices. - rest_lengths: NDArray[np.floating] + rest_lengths: NDArray[np.float64] 1D (n_elems) array containing data with 'float' type. Rod element lengths at rest configuration. - density: NDArray[np.floating] + density: NDArray[np.float64] 1D (n_elems) array containing data with 'float' type. Rod elements densities. - volume: NDArray[np.floating] + volume: NDArray[np.float64] 1D (n_elems) array containing data with 'float' type. Rod element volumes. - mass: NDArray[np.floating] + mass: NDArray[np.float64] 1D (n_nodes) array containing data with 'float' type. Rod node masses. Note that masses are stored on the nodes, not on elements. - mass_second_moment_of_inertia: NDArray[np.floating] + mass_second_moment_of_inertia: NDArray[np.float64] 3D (dim, dim, n_elems) array containing data with 'float' type. Rod element mass second moment of interia. - inv_mass_second_moment_of_inertia: NDArray[np.floating] + inv_mass_second_moment_of_inertia: NDArray[np.float64] 3D (dim, dim, n_elems) array containing data with 'float' type. Rod element inverse mass moment of inertia. - rest_voronoi_lengths: NDArray[np.floating] + rest_voronoi_lengths: NDArray[np.float64] 1D (n_voronoi) array containing data with 'float' type. Rod lengths on the voronoi domain at the rest configuration. - internal_forces: NDArray[np.floating] + internal_forces: NDArray[np.float64] 2D (dim, n_nodes) array containing data with 'float' type. Rod node internal forces. Note that internal forces are stored on the node, not on elements. - internal_torques: NDArray[np.floating] + internal_torques: NDArray[np.float64] 2D (dim, n_elems) array containing data with 'float' type. Rod element internal torques. - external_forces: NDArray[np.floating] + external_forces: NDArray[np.float64] 2D (dim, n_nodes) array containing data with 'float' type. External forces acting on rod nodes. - external_torques: NDArray[np.floating] + external_torques: NDArray[np.float64] 2D (dim, n_elems) array containing data with 'float' type. External torques acting on rod elements. - lengths: NDArray[np.floating] + lengths: NDArray[np.float64] 1D (n_elems) array containing data with 'float' type. Rod element lengths. - tangents: NDArray[np.floating] + tangents: NDArray[np.float64] 2D (dim, n_elems) array containing data with 'float' type. Rod element tangent vectors. - radius: NDArray[np.floating] + radius: NDArray[np.float64] 1D (n_elems) array containing data with 'float' type. Rod element radius. - dilatation: NDArray[np.floating] + dilatation: NDArray[np.float64] 1D (n_elems) array containing data with 'float' type. Rod element dilatation. - voronoi_dilatation: NDArray[np.floating] + voronoi_dilatation: NDArray[np.float64] 1D (n_voronoi) array containing data with 'float' type. Rod dilatation on voronoi domain. - dilatation_rate: NDArray[np.floating] + dilatation_rate: NDArray[np.float64] 1D (n_elems) array containing data with 'float' type. Rod element dilatation rates. """ @@ -155,37 +155,37 @@ class CosseratRod(RodBase, KnotTheory): def __init__( self: CosseratRodProtocol, n_elements: int, - position: NDArray[np.floating], - velocity: NDArray[np.floating], - omega: NDArray[np.floating], - acceleration: NDArray[np.floating], - angular_acceleration: NDArray[np.floating], - directors: NDArray[np.floating], - radius: NDArray[np.floating], - mass_second_moment_of_inertia: NDArray[np.floating], - inv_mass_second_moment_of_inertia: NDArray[np.floating], - shear_matrix: NDArray[np.floating], - bend_matrix: NDArray[np.floating], - density_array: NDArray[np.floating], - volume: NDArray[np.floating], - mass: NDArray[np.floating], - internal_forces: NDArray[np.floating], - internal_torques: NDArray[np.floating], - external_forces: NDArray[np.floating], - external_torques: NDArray[np.floating], - lengths: NDArray[np.floating], - rest_lengths: NDArray[np.floating], - tangents: NDArray[np.floating], - dilatation: NDArray[np.floating], - dilatation_rate: NDArray[np.floating], - voronoi_dilatation: NDArray[np.floating], - rest_voronoi_lengths: NDArray[np.floating], - sigma: NDArray[np.floating], - kappa: NDArray[np.floating], - rest_sigma: NDArray[np.floating], - rest_kappa: NDArray[np.floating], - internal_stress: NDArray[np.floating], - internal_couple: NDArray[np.floating], + position: NDArray[np.float64], + velocity: NDArray[np.float64], + omega: NDArray[np.float64], + acceleration: NDArray[np.float64], + angular_acceleration: NDArray[np.float64], + directors: NDArray[np.float64], + radius: NDArray[np.float64], + mass_second_moment_of_inertia: NDArray[np.float64], + inv_mass_second_moment_of_inertia: NDArray[np.float64], + shear_matrix: NDArray[np.float64], + bend_matrix: NDArray[np.float64], + density_array: NDArray[np.float64], + volume: NDArray[np.float64], + mass: NDArray[np.float64], + internal_forces: NDArray[np.float64], + internal_torques: NDArray[np.float64], + external_forces: NDArray[np.float64], + external_torques: NDArray[np.float64], + lengths: NDArray[np.float64], + rest_lengths: NDArray[np.float64], + tangents: NDArray[np.float64], + dilatation: NDArray[np.float64], + dilatation_rate: NDArray[np.float64], + voronoi_dilatation: NDArray[np.float64], + rest_voronoi_lengths: NDArray[np.float64], + sigma: NDArray[np.float64], + kappa: NDArray[np.float64], + rest_sigma: NDArray[np.float64], + rest_kappa: NDArray[np.float64], + internal_stress: NDArray[np.float64], + internal_couple: NDArray[np.float64], ring_rod_flag: bool, ) -> None: self.n_nodes = n_elements + 1 if ring_rod_flag else n_elements @@ -250,14 +250,14 @@ def __init__( def straight_rod( cls, n_elements: int, - start: NDArray[np.floating], - direction: NDArray[np.floating], - normal: NDArray[np.floating], + start: NDArray[np.float64], + direction: NDArray[np.float64], + normal: NDArray[np.float64], base_length: float, base_radius: float, density: float, *, - nu: Optional[np.floating] = None, + nu: Optional[np.float64] = None, youngs_modulus: float, **kwargs: Any, ) -> Self: @@ -276,11 +276,11 @@ def straight_rod( n_elements : int Number of element. Must be greater than 3. Generally recommended to start with 40-50, and adjust the resolution. - start : NDArray[np.floating] + start : NDArray[np.float64] Starting coordinate in 3D - direction : NDArray[np.floating] + direction : NDArray[np.float64] Direction of the rod in 3D - normal : NDArray[np.floating] + normal : NDArray[np.float64] Normal vector of the rod in 3D base_length : float Total length of the rod @@ -398,9 +398,9 @@ def straight_rod( def ring_rod( cls, n_elements: int, - ring_center_position: NDArray[np.floating], - direction: NDArray[np.floating], - normal: NDArray[np.floating], + ring_center_position: NDArray[np.float64], + direction: NDArray[np.float64], + normal: NDArray[np.float64], base_length: float, base_radius: float, density: float, @@ -423,11 +423,11 @@ def ring_rod( ---------- n_elements : int Number of element. Must be greater than 3. Generarally recommended to start with 40-50, and adjust the resolution. - ring_center_position : NDArray[np.floating] + ring_center_position : NDArray[np.float64] Center coordinate for ring rod in 3D - direction : NDArray[np.floating] + direction : NDArray[np.float64] Direction of the rod in 3D - normal : NDArray[np.floating] + normal : NDArray[np.float64] Normal vector of the rod in 3D base_length : float Total length of the rod @@ -545,7 +545,7 @@ def ring_rod( return rod def compute_internal_forces_and_torques( - self: CosseratRodProtocol, time: np.floating + self: CosseratRodProtocol, time: np.float64 ) -> None: """ Compute internal forces and torques. We need to compute internal forces and torques before the acceleration because @@ -555,7 +555,7 @@ def compute_internal_forces_and_torques( Parameters ---------- - time: np.floating + time: np.float64 current time """ @@ -601,13 +601,13 @@ def compute_internal_forces_and_torques( ) # Interface to time-stepper mixins (Symplectic, Explicit), which calls this method - def update_accelerations(self: CosseratRodProtocol, time: np.floating) -> None: + def update_accelerations(self: CosseratRodProtocol, time: np.float64) -> None: """ Updates the acceleration variables Parameters ---------- - time: np.floating + time: np.float64 current time """ @@ -624,13 +624,13 @@ def update_accelerations(self: CosseratRodProtocol, time: np.floating) -> None: ) def zeroed_out_external_forces_and_torques( - self: CosseratRodProtocol, time: np.floating + self: CosseratRodProtocol, time: np.float64 ) -> None: _zeroed_out_external_forces_and_torques( self.external_forces, self.external_torques ) - def compute_translational_energy(self: CosseratRodProtocol) -> NDArray[np.floating]: + def compute_translational_energy(self: CosseratRodProtocol) -> NDArray[np.float64]: """ Compute total translational energy of the rod at the instance. """ @@ -644,7 +644,7 @@ def compute_translational_energy(self: CosseratRodProtocol) -> NDArray[np.floati ).sum() ) - def compute_rotational_energy(self: CosseratRodProtocol) -> NDArray[np.floating]: + def compute_rotational_energy(self: CosseratRodProtocol) -> NDArray[np.float64]: """ Compute total rotational energy of the rod at the instance. """ @@ -656,7 +656,7 @@ def compute_rotational_energy(self: CosseratRodProtocol) -> NDArray[np.floating] def compute_velocity_center_of_mass( self: CosseratRodProtocol, - ) -> NDArray[np.floating]: + ) -> NDArray[np.float64]: """ Compute velocity center of mass of the rod at the instance. """ @@ -667,7 +667,7 @@ def compute_velocity_center_of_mass( def compute_position_center_of_mass( self: CosseratRodProtocol, - ) -> NDArray[np.floating]: + ) -> NDArray[np.float64]: """ Compute position center of mass of the rod at the instance. """ @@ -676,7 +676,7 @@ def compute_position_center_of_mass( return sum_mass_times_position / self.mass.sum() - def compute_bending_energy(self: CosseratRodProtocol) -> NDArray[np.floating]: + def compute_bending_energy(self: CosseratRodProtocol) -> NDArray[np.float64]: """ Compute total bending energy of the rod at the instance. """ @@ -692,7 +692,7 @@ def compute_bending_energy(self: CosseratRodProtocol) -> NDArray[np.floating]: ).sum() ) - def compute_shear_energy(self: CosseratRodProtocol) -> NDArray[np.floating]: + def compute_shear_energy(self: CosseratRodProtocol) -> NDArray[np.float64]: """ Compute total shear energy of the rod at the instance. """ @@ -711,11 +711,11 @@ def compute_shear_energy(self: CosseratRodProtocol) -> NDArray[np.floating]: @numba.njit(cache=True) # type: ignore def _compute_geometry_from_state( - position_collection: NDArray[np.floating], - volume: NDArray[np.floating], - lengths: NDArray[np.floating], - tangents: NDArray[np.floating], - radius: NDArray[np.floating], + position_collection: NDArray[np.float64], + volume: NDArray[np.float64], + lengths: NDArray[np.float64], + tangents: NDArray[np.float64], + radius: NDArray[np.float64], ) -> None: """ Update given . @@ -739,15 +739,15 @@ def _compute_geometry_from_state( @numba.njit(cache=True) # type: ignore def _compute_all_dilatations( - position_collection: NDArray[np.floating], - volume: NDArray[np.floating], - lengths: NDArray[np.floating], - tangents: NDArray[np.floating], - radius: NDArray[np.floating], - dilatation: NDArray[np.floating], - rest_lengths: NDArray[np.floating], - rest_voronoi_lengths: NDArray[np.floating], - voronoi_dilatation: NDArray[np.floating], + position_collection: NDArray[np.float64], + volume: NDArray[np.float64], + lengths: NDArray[np.float64], + tangents: NDArray[np.float64], + radius: NDArray[np.float64], + dilatation: NDArray[np.float64], + rest_lengths: NDArray[np.float64], + rest_voronoi_lengths: NDArray[np.float64], + voronoi_dilatation: NDArray[np.float64], ) -> None: """ Update @@ -769,11 +769,11 @@ def _compute_all_dilatations( @numba.njit(cache=True) # type: ignore def _compute_dilatation_rate( - position_collection: NDArray[np.floating], - velocity_collection: NDArray[np.floating], - lengths: NDArray[np.floating], - rest_lengths: NDArray[np.floating], - dilatation_rate: NDArray[np.floating], + position_collection: NDArray[np.float64], + velocity_collection: NDArray[np.float64], + lengths: NDArray[np.float64], + rest_lengths: NDArray[np.float64], + dilatation_rate: NDArray[np.float64], ) -> None: """ Update dilatation_rate given position, velocity, length, and rest_length @@ -800,17 +800,17 @@ def _compute_dilatation_rate( @numba.njit(cache=True) # type: ignore def _compute_shear_stretch_strains( - position_collection: NDArray[np.floating], - volume: NDArray[np.floating], - lengths: NDArray[np.floating], - tangents: NDArray[np.floating], - radius: NDArray[np.floating], - rest_lengths: NDArray[np.floating], - rest_voronoi_lengths: NDArray[np.floating], - dilatation: NDArray[np.floating], - voronoi_dilatation: NDArray[np.floating], - director_collection: NDArray[np.floating], - sigma: NDArray[np.floating], + position_collection: NDArray[np.float64], + volume: NDArray[np.float64], + lengths: NDArray[np.float64], + tangents: NDArray[np.float64], + radius: NDArray[np.float64], + rest_lengths: NDArray[np.float64], + rest_voronoi_lengths: NDArray[np.float64], + dilatation: NDArray[np.float64], + voronoi_dilatation: NDArray[np.float64], + director_collection: NDArray[np.float64], + sigma: NDArray[np.float64], ) -> None: """ Update given . @@ -835,20 +835,20 @@ def _compute_shear_stretch_strains( @numba.njit(cache=True) # type: ignore def _compute_internal_shear_stretch_stresses_from_model( - position_collection: NDArray[np.floating], - volume: NDArray[np.floating], - lengths: NDArray[np.floating], - tangents: NDArray[np.floating], - radius: NDArray[np.floating], - rest_lengths: NDArray[np.floating], - rest_voronoi_lengths: NDArray[np.floating], - dilatation: NDArray[np.floating], - voronoi_dilatation: NDArray[np.floating], - director_collection: NDArray[np.floating], - sigma: NDArray[np.floating], - rest_sigma: NDArray[np.floating], - shear_matrix: NDArray[np.floating], - internal_stress: NDArray[np.floating], + position_collection: NDArray[np.float64], + volume: NDArray[np.float64], + lengths: NDArray[np.float64], + tangents: NDArray[np.float64], + radius: NDArray[np.float64], + rest_lengths: NDArray[np.float64], + rest_voronoi_lengths: NDArray[np.float64], + dilatation: NDArray[np.float64], + voronoi_dilatation: NDArray[np.float64], + director_collection: NDArray[np.float64], + sigma: NDArray[np.float64], + rest_sigma: NDArray[np.float64], + shear_matrix: NDArray[np.float64], + internal_stress: NDArray[np.float64], ) -> None: """ Update given . @@ -875,9 +875,9 @@ def _compute_internal_shear_stretch_stresses_from_model( @numba.njit(cache=True) # type: ignore def _compute_bending_twist_strains( - director_collection: NDArray[np.floating], - rest_voronoi_lengths: NDArray[np.floating], - kappa: NDArray[np.floating], + director_collection: NDArray[np.float64], + rest_voronoi_lengths: NDArray[np.float64], + kappa: NDArray[np.float64], ) -> None: """ Update given . @@ -892,12 +892,12 @@ def _compute_bending_twist_strains( @numba.njit(cache=True) # type: ignore def _compute_internal_bending_twist_stresses_from_model( - director_collection: NDArray[np.floating], - rest_voronoi_lengths: NDArray[np.floating], - internal_couple: NDArray[np.floating], - bend_matrix: NDArray[np.floating], - kappa: NDArray[np.floating], - rest_kappa: NDArray[np.floating], + director_collection: NDArray[np.float64], + rest_voronoi_lengths: NDArray[np.float64], + internal_couple: NDArray[np.float64], + bend_matrix: NDArray[np.float64], + kappa: NDArray[np.float64], + rest_kappa: NDArray[np.float64], ) -> None: """ Upate given . @@ -921,22 +921,22 @@ def _compute_internal_bending_twist_stresses_from_model( @numba.njit(cache=True) # type: ignore def _compute_internal_forces( - position_collection: NDArray[np.floating], - volume: NDArray[np.floating], - lengths: NDArray[np.floating], - tangents: NDArray[np.floating], - radius: NDArray[np.floating], - rest_lengths: NDArray[np.floating], - rest_voronoi_lengths: NDArray[np.floating], - dilatation: NDArray[np.floating], - voronoi_dilatation: NDArray[np.floating], - director_collection: NDArray[np.floating], - sigma: NDArray[np.floating], - rest_sigma: NDArray[np.floating], - shear_matrix: NDArray[np.floating], - internal_stress: NDArray[np.floating], - internal_forces: NDArray[np.floating], - ghost_elems_idx: NDArray[np.floating], + position_collection: NDArray[np.float64], + volume: NDArray[np.float64], + lengths: NDArray[np.float64], + tangents: NDArray[np.float64], + radius: NDArray[np.float64], + rest_lengths: NDArray[np.float64], + rest_voronoi_lengths: NDArray[np.float64], + dilatation: NDArray[np.float64], + voronoi_dilatation: NDArray[np.float64], + director_collection: NDArray[np.float64], + sigma: NDArray[np.float64], + rest_sigma: NDArray[np.float64], + shear_matrix: NDArray[np.float64], + internal_stress: NDArray[np.float64], + internal_forces: NDArray[np.float64], + ghost_elems_idx: NDArray[np.float64], ) -> None: """ Update given . @@ -982,25 +982,25 @@ def _compute_internal_forces( @numba.njit(cache=True) # type: ignore def _compute_internal_torques( - position_collection: NDArray[np.floating], - velocity_collection: NDArray[np.floating], - tangents: NDArray[np.floating], - lengths: NDArray[np.floating], - rest_lengths: NDArray[np.floating], - director_collection: NDArray[np.floating], - rest_voronoi_lengths: NDArray[np.floating], - bend_matrix: NDArray[np.floating], - rest_kappa: NDArray[np.floating], - kappa: NDArray[np.floating], - voronoi_dilatation: NDArray[np.floating], - mass_second_moment_of_inertia: NDArray[np.floating], - omega_collection: NDArray[np.floating], - internal_stress: NDArray[np.floating], - internal_couple: NDArray[np.floating], - dilatation: NDArray[np.floating], - dilatation_rate: NDArray[np.floating], - internal_torques: NDArray[np.floating], - ghost_voronoi_idx: NDArray[np.integer], + position_collection: NDArray[np.float64], + velocity_collection: NDArray[np.float64], + tangents: NDArray[np.float64], + lengths: NDArray[np.float64], + rest_lengths: NDArray[np.float64], + director_collection: NDArray[np.float64], + rest_voronoi_lengths: NDArray[np.float64], + bend_matrix: NDArray[np.float64], + rest_kappa: NDArray[np.float64], + kappa: NDArray[np.float64], + voronoi_dilatation: NDArray[np.float64], + mass_second_moment_of_inertia: NDArray[np.float64], + omega_collection: NDArray[np.float64], + internal_stress: NDArray[np.float64], + internal_couple: NDArray[np.float64], + dilatation: NDArray[np.float64], + dilatation_rate: NDArray[np.float64], + internal_torques: NDArray[np.float64], + ghost_voronoi_idx: NDArray[np.int32], ) -> None: """ Update . @@ -1071,15 +1071,15 @@ def _compute_internal_torques( @numba.njit(cache=True) # type: ignore def _update_accelerations( - acceleration_collection: NDArray[np.floating], - internal_forces: NDArray[np.floating], - external_forces: NDArray[np.floating], - mass: NDArray[np.floating], - alpha_collection: NDArray[np.floating], - inv_mass_second_moment_of_inertia: NDArray[np.floating], - internal_torques: NDArray[np.floating], - external_torques: NDArray[np.floating], - dilatation: NDArray[np.floating], + acceleration_collection: NDArray[np.float64], + internal_forces: NDArray[np.float64], + external_forces: NDArray[np.float64], + mass: NDArray[np.float64], + alpha_collection: NDArray[np.float64], + inv_mass_second_moment_of_inertia: NDArray[np.float64], + internal_torques: NDArray[np.float64], + external_torques: NDArray[np.float64], + dilatation: NDArray[np.float64], ) -> None: """ Update given . @@ -1106,7 +1106,7 @@ def _update_accelerations( @numba.njit(cache=True) # type: ignore def _zeroed_out_external_forces_and_torques( - external_forces: NDArray[np.floating], external_torques: NDArray[np.floating] + external_forces: NDArray[np.float64], external_torques: NDArray[np.float64] ) -> None: """ This function is to zeroed out external forces and torques. diff --git a/elastica/rod/data_structures.py b/elastica/rod/data_structures.py index 46588785b..f7b248bcc 100644 --- a/elastica/rod/data_structures.py +++ b/elastica/rod/data_structures.py @@ -69,9 +69,9 @@ def __init__(self: SymplecticSystemProtocol) -> None: def dynamic_rates( self: SymplecticSystemProtocol, - time: np.floating, - prefac: np.floating, - ) -> NDArray[np.floating]: + time: np.float64, + prefac: np.float64, + ) -> NDArray[np.float64]: self.update_accelerations(time) return self.dynamic_states.dynamic_rates(time, prefac) @@ -79,18 +79,18 @@ def dynamic_rates( def _bootstrap_from_data( stepper_type: str, n_elems: int, - vector_states: NDArray[np.floating], - matrix_states: NDArray[np.floating], + vector_states: NDArray[np.float64], + matrix_states: NDArray[np.float64], ) -> Optional[ tuple[ "_State", "_DerivativeState", - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], ] ]: """Returns states wrapping numpy arrays based on the time-stepping algorithm @@ -175,9 +175,9 @@ class _State: def __init__( self, n_elems: int, - position_collection_view: NDArray[np.floating], - director_collection_view: NDArray[np.floating], - kinematic_rate_collection_view: NDArray[np.floating], + position_collection_view: NDArray[np.float64], + director_collection_view: NDArray[np.float64], + kinematic_rate_collection_view: NDArray[np.float64], ) -> None: """ Parameters @@ -194,7 +194,7 @@ def __init__( self.director_collection = director_collection_view self.kinematic_rate_collection = kinematic_rate_collection_view - def __iadd__(self, scaled_deriv_array: NDArray[np.floating]) -> Self: + def __iadd__(self, scaled_deriv_array: NDArray[np.float64]) -> Self: """overloaded += operator The add for directors is customized to reflect Rodrigues' rotation @@ -263,7 +263,7 @@ def __iadd__(self, scaled_deriv_array: NDArray[np.floating]) -> Self: ] return self - def __add__(self, scaled_derivative_state: NDArray[np.floating]) -> "_State": + def __add__(self, scaled_derivative_state: NDArray[np.float64]) -> "_State": """overloaded + operator, useful in state.k1 = state + dt * deriv_state The add for directors is customized to reflect Rodrigues' rotation @@ -318,7 +318,7 @@ class _DerivativeState: """ def __init__( - self, _unused_n_elems: int, rate_collection_view: NDArray[np.floating] + self, _unused_n_elems: int, rate_collection_view: NDArray[np.float64] ) -> None: """ Parameters @@ -330,7 +330,7 @@ def __init__( super(_DerivativeState, self).__init__() self.rate_collection = rate_collection_view - def __rmul__(self, scalar: np.floating) -> NDArray[np.floating]: # type: ignore + def __rmul__(self, scalar: np.float64) -> NDArray[np.float64]: # type: ignore """overloaded scalar * self, Parameters @@ -378,7 +378,7 @@ def __rmul__(self, scalar: np.floating) -> NDArray[np.floating]: # type: ignore """ return scalar * self.rate_collection - def __mul__(self, scalar: np.floating) -> NDArray[np.floating]: + def __mul__(self, scalar: np.float64) -> NDArray[np.float64]: """overloaded self * scalar TODO Check if this pattern (forwarding to __mul__) has @@ -413,8 +413,8 @@ class _KinematicState: def __init__( self, - position_collection_view: NDArray[np.floating], - director_collection_view: NDArray[np.floating], + position_collection_view: NDArray[np.float64], + director_collection_view: NDArray[np.float64], ) -> None: """ Parameters @@ -431,11 +431,11 @@ def __init__( @njit(cache=True) # type: ignore def overload_operator_kinematic_numba( n_nodes: int, - prefac: np.floating, - position_collection: NDArray[np.floating], - director_collection: NDArray[np.floating], - velocity_collection: NDArray[np.floating], - omega_collection: NDArray[np.floating], + prefac: np.float64, + position_collection: NDArray[np.float64], + director_collection: NDArray[np.float64], + velocity_collection: NDArray[np.float64], + omega_collection: NDArray[np.float64], ) -> None: """overloaded += operator @@ -476,10 +476,10 @@ class _DynamicState: def __init__( self, - v_w_collection: NDArray[np.floating], - dvdt_dwdt_collection: NDArray[np.floating], - velocity_collection: NDArray[np.floating], - omega_collection: NDArray[np.floating], + v_w_collection: NDArray[np.float64], + dvdt_dwdt_collection: NDArray[np.float64], + velocity_collection: NDArray[np.float64], + omega_collection: NDArray[np.float64], ) -> None: """ Parameters @@ -498,8 +498,8 @@ def __init__( self.omega_collection = omega_collection def kinematic_rates( - self, time: np.floating, prefac: np.floating - ) -> tuple[NDArray[np.floating], NDArray[np.floating]]: + self, time: np.float64, prefac: np.float64 + ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: """Yields kinematic rates to interact with _KinematicState Returns @@ -516,8 +516,8 @@ def kinematic_rates( return self.velocity_collection, self.omega_collection def dynamic_rates( - self, time: np.floating, prefac: np.floating - ) -> NDArray[np.floating]: + self, time: np.float64, prefac: np.float64 + ) -> NDArray[np.float64]: """Yields dynamic rates to add to with _DynamicState Returns ------- @@ -533,8 +533,8 @@ def dynamic_rates( @njit(cache=True) # type: ignore def overload_operator_dynamic_numba( - rate_collection: NDArray[np.floating], - scaled_second_deriv_array: NDArray[np.floating], + rate_collection: NDArray[np.float64], + scaled_second_deriv_array: NDArray[np.float64], ) -> None: """overloaded += operator, updating dynamic_rates Parameters diff --git a/elastica/rod/factory_function.py b/elastica/rod/factory_function.py index 22cc33bc6..2765e8726 100644 --- a/elastica/rod/factory_function.py +++ b/elastica/rod/factory_function.py @@ -10,16 +10,16 @@ def allocate( n_elements: int, - direction: NDArray[np.floating], - normal: NDArray[np.floating], - base_length: np.floating, - base_radius: np.floating, - density: np.floating, - youngs_modulus: np.floating, + direction: NDArray[np.float64], + normal: NDArray[np.float64], + base_length: np.float64, + base_radius: np.float64, + density: np.float64, + youngs_modulus: np.float64, *, rod_origin_position: np.ndarray, ring_rod_flag: bool, - shear_modulus: Optional[np.floating] = None, + shear_modulus: Optional[np.float64] = None, position: Optional[np.ndarray] = None, directors: Optional[np.ndarray] = None, rest_sigma: Optional[np.ndarray] = None, @@ -27,37 +27,37 @@ def allocate( **kwargs: Any, ) -> tuple[ int, - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], - NDArray[np.floating], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], + NDArray[np.float64], ]: log = logging.getLogger() @@ -388,7 +388,7 @@ def _assert_shape( def _position_validity_checker( - position: NDArray[np.floating], start: NDArray[np.floating], n_elements: int + position: NDArray[np.float64], start: NDArray[np.float64], n_elements: int ) -> None: """Checker on user-defined position validity""" _assert_shape(position, (MaxDimension.value(), n_elements + 1), "position") @@ -406,7 +406,7 @@ def _position_validity_checker( def _directors_validity_checker( - directors: NDArray[np.floating], tangents: NDArray[np.floating], n_elements: int + directors: NDArray[np.float64], tangents: NDArray[np.float64], n_elements: int ) -> None: """Checker on user-defined directors validity""" _assert_shape( @@ -454,8 +454,8 @@ def _directors_validity_checker( def _position_validity_checker_ring_rod( - position: NDArray[np.floating], - ring_center_position: NDArray[np.floating], + position: NDArray[np.float64], + ring_center_position: NDArray[np.float64], n_elements: int, ) -> None: """Checker on user-defined position validity""" diff --git a/elastica/rod/knot_theory.py b/elastica/rod/knot_theory.py index ac440d140..b5f00184c 100644 --- a/elastica/rod/knot_theory.py +++ b/elastica/rod/knot_theory.py @@ -58,7 +58,7 @@ def __init__(self) -> None: """ - def compute_twist(self: CosseratRodProtocol) -> NDArray[np.floating]: + def compute_twist(self: CosseratRodProtocol) -> NDArray[np.float64]: """ See :ref:`api/rods:Knot Theory (Mixin)` for the detail. """ @@ -72,7 +72,7 @@ def compute_writhe( self: CosseratRodProtocol, type_of_additional_segment: str = "next_tangent", alpha: float = 1.0, - ) -> NDArray[np.floating]: + ) -> NDArray[np.float64]: """ See :ref:`api/rods:Knot Theory (Mixin)` for the detail. @@ -95,7 +95,7 @@ def compute_link( self: CosseratRodProtocol, type_of_additional_segment: str = "next_tangent", alpha: float = 1.0, - ) -> NDArray[np.floating]: + ) -> NDArray[np.float64]: """ See :ref:`api/rods:Knot Theory (Mixin)` for the detail. @@ -119,8 +119,8 @@ def compute_link( def compute_twist( - center_line: NDArray[np.floating], normal_collection: NDArray[np.floating] -) -> tuple[NDArray[np.floating], NDArray[np.floating]]: + center_line: NDArray[np.float64], normal_collection: NDArray[np.float64] +) -> tuple[NDArray[np.float64], NDArray[np.float64]]: """ Compute the twist of a rod, using center_line and normal collection. @@ -172,8 +172,8 @@ def compute_twist( @njit(cache=True) # type: ignore def _compute_twist( - center_line: NDArray[np.floating], normal_collection: NDArray[np.floating] -) -> tuple[NDArray[np.floating], NDArray[np.floating]]: + center_line: NDArray[np.float64], normal_collection: NDArray[np.float64] +) -> tuple[NDArray[np.float64], NDArray[np.float64]]: """ Parameters ---------- @@ -249,10 +249,10 @@ def _compute_twist( def compute_writhe( - center_line: NDArray[np.floating], - segment_length: np.floating, + center_line: NDArray[np.float64], + segment_length: np.float64, type_of_additional_segment: str, -) -> NDArray[np.floating]: +) -> NDArray[np.float64]: """ This function computes the total writhe history of a rod. @@ -302,7 +302,7 @@ def compute_writhe( @njit(cache=True) # type: ignore -def _compute_writhe(center_line: NDArray[np.floating]) -> NDArray[np.floating]: +def _compute_writhe(center_line: NDArray[np.float64]) -> NDArray[np.float64]: """ Parameters ---------- @@ -374,12 +374,12 @@ def _compute_writhe(center_line: NDArray[np.floating]) -> NDArray[np.floating]: def compute_link( - center_line: NDArray[np.floating], - normal_collection: NDArray[np.floating], - radius: NDArray[np.floating], - segment_length: np.floating, + center_line: NDArray[np.float64], + normal_collection: NDArray[np.float64], + radius: NDArray[np.float64], + segment_length: np.float64, type_of_additional_segment: str, -) -> NDArray[np.floating]: +) -> NDArray[np.float64]: """ This function computes the total link history of a rod. @@ -459,10 +459,10 @@ def compute_link( @njit(cache=True) # type: ignore def _compute_auxiliary_line( - center_line: NDArray[np.floating], - normal_collection: NDArray[np.floating], - radius: NDArray[np.floating], -) -> NDArray[np.floating]: + center_line: NDArray[np.float64], + normal_collection: NDArray[np.float64], + radius: NDArray[np.float64], +) -> NDArray[np.float64]: """ This function computes the auxiliary line using rod center line and normal collection. @@ -518,8 +518,8 @@ def _compute_auxiliary_line( @njit(cache=True) # type: ignore def _compute_link( - center_line: NDArray[np.floating], auxiliary_line: NDArray[np.floating] -) -> NDArray[np.floating]: + center_line: NDArray[np.float64], auxiliary_line: NDArray[np.float64] +) -> NDArray[np.float64]: """ Parameters @@ -598,11 +598,11 @@ def _compute_link( @njit(cache=True) # type: ignore def _compute_auxiliary_line_added_segments( - beginning_direction: NDArray[np.floating], - end_direction: NDArray[np.floating], - auxiliary_line: NDArray[np.floating], - segment_length: np.floating, -) -> NDArray[np.floating]: + beginning_direction: NDArray[np.float64], + end_direction: NDArray[np.float64], + auxiliary_line: NDArray[np.float64], + segment_length: np.float64, +) -> NDArray[np.float64]: """ This code is for computing position of added segments to the auxiliary line. @@ -644,10 +644,10 @@ def _compute_auxiliary_line_added_segments( @njit(cache=True) # type: ignore def _compute_additional_segment( - center_line: NDArray[np.floating], - segment_length: np.floating, + center_line: NDArray[np.float64], + segment_length: np.float64, type_of_additional_segment: str, -) -> tuple[NDArray[np.floating], NDArray[np.floating], NDArray[np.floating]]: +) -> tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]]: """ This function adds two points at the end of center line. Distance from the center line is given by segment_length. Direction from center line to the new point locations can be computed using 3 methods, which can be selected by diff --git a/elastica/rod/protocol.py b/elastica/rod/protocol.py index 62d33fe57..c28de3bbb 100644 --- a/elastica/rod/protocol.py +++ b/elastica/rod/protocol.py @@ -16,34 +16,34 @@ class CosseratRodProtocol( SystemProtocol, SlenderBodyGeometryProtocol, _RodEnergy, Protocol ): - mass: NDArray[np.floating] - volume: NDArray[np.floating] - radius: NDArray[np.floating] - tangents: NDArray[np.floating] - lengths: NDArray[np.floating] - rest_lengths: NDArray[np.floating] - rest_voronoi_lengths: NDArray[np.floating] - kappa: NDArray[np.floating] - sigma: NDArray[np.floating] - rest_kappa: NDArray[np.floating] - rest_sigma: NDArray[np.floating] - - internal_stress: NDArray[np.floating] - internal_couple: NDArray[np.floating] - dilatation: NDArray[np.floating] - dilatation_rate: NDArray[np.floating] - voronoi_dilatation: NDArray[np.floating] - - bend_matrix: NDArray[np.floating] - shear_matrix: NDArray[np.floating] - - mass_second_moment_of_inertia: NDArray[np.floating] - inv_mass_second_moment_of_inertia: NDArray[np.floating] - - ghost_voronoi_idx: NDArray[np.integer] - ghost_elems_idx: NDArray[np.integer] + mass: NDArray[np.float64] + volume: NDArray[np.float64] + radius: NDArray[np.float64] + tangents: NDArray[np.float64] + lengths: NDArray[np.float64] + rest_lengths: NDArray[np.float64] + rest_voronoi_lengths: NDArray[np.float64] + kappa: NDArray[np.float64] + sigma: NDArray[np.float64] + rest_kappa: NDArray[np.float64] + rest_sigma: NDArray[np.float64] + + internal_stress: NDArray[np.float64] + internal_couple: NDArray[np.float64] + dilatation: NDArray[np.float64] + dilatation_rate: NDArray[np.float64] + voronoi_dilatation: NDArray[np.float64] + + bend_matrix: NDArray[np.float64] + shear_matrix: NDArray[np.float64] + + mass_second_moment_of_inertia: NDArray[np.float64] + inv_mass_second_moment_of_inertia: NDArray[np.float64] + + ghost_voronoi_idx: NDArray[np.int32] + ghost_elems_idx: NDArray[np.int32] ring_rod_flag: bool - periodic_boundary_nodes_idx: NDArray[np.integer] - periodic_boundary_elems_idx: NDArray[np.integer] - periodic_boundary_voronoi_idx: NDArray[np.integer] + periodic_boundary_nodes_idx: NDArray[np.int32] + periodic_boundary_elems_idx: NDArray[np.int32] + periodic_boundary_voronoi_idx: NDArray[np.int32] diff --git a/elastica/rod/rod_base.py b/elastica/rod/rod_base.py index 11dc9211f..3bd846b7b 100644 --- a/elastica/rod/rod_base.py +++ b/elastica/rod/rod_base.py @@ -21,18 +21,18 @@ def __init__(self) -> None: """ RodBase does not take any arguments. """ - self.position_collection: NDArray[np.floating] - self.velocity_collection: NDArray[np.floating] - self.acceleration_collection: NDArray[np.floating] - self.director_collection: NDArray[np.floating] - self.omega_collection: NDArray[np.floating] - self.alpha_collection: NDArray[np.floating] - self.external_forces: NDArray[np.floating] - self.external_torques: NDArray[np.floating] - - self.ghost_voronoi_idx: NDArray[np.integer] - self.ghost_elems_idx: NDArray[np.integer] - - self.periodic_boundary_nodes_idx: NDArray[np.integer] - self.periodic_boundary_elems_idx: NDArray[np.integer] - self.periodic_boundary_voronoi_idx: NDArray[np.integer] + self.position_collection: NDArray[np.float64] + self.velocity_collection: NDArray[np.float64] + self.acceleration_collection: NDArray[np.float64] + self.director_collection: NDArray[np.float64] + self.omega_collection: NDArray[np.float64] + self.alpha_collection: NDArray[np.float64] + self.external_forces: NDArray[np.float64] + self.external_torques: NDArray[np.float64] + + self.ghost_voronoi_idx: NDArray[np.int32] + self.ghost_elems_idx: NDArray[np.int32] + + self.periodic_boundary_nodes_idx: NDArray[np.int32] + self.periodic_boundary_elems_idx: NDArray[np.int32] + self.periodic_boundary_voronoi_idx: NDArray[np.int32] diff --git a/elastica/surface/plane.py b/elastica/surface/plane.py index 99b8525d8..b12a3feaf 100644 --- a/elastica/surface/plane.py +++ b/elastica/surface/plane.py @@ -8,7 +8,7 @@ class Plane(SurfaceBase): def __init__( - self, plane_origin: NDArray[np.floating], plane_normal: NDArray[np.floating] + self, plane_origin: NDArray[np.float64], plane_normal: NDArray[np.float64] ): """ Plane surface initializer. diff --git a/elastica/surface/surface_base.py b/elastica/surface/surface_base.py index 9870afe96..37812ea7d 100644 --- a/elastica/surface/surface_base.py +++ b/elastica/surface/surface_base.py @@ -21,8 +21,8 @@ def __init__(self) -> None: """ SurfaceBase does not take any arguments. """ - self.normal: NDArray[np.floating] # (3,) - self.origin: NDArray[np.floating] # (3, 1) + self.normal: NDArray[np.float64] # (3,) + self.origin: NDArray[np.float64] # (3, 1) if TYPE_CHECKING: diff --git a/elastica/systems/protocol.py b/elastica/systems/protocol.py index ccac73f64..254cdcaf9 100644 --- a/elastica/systems/protocol.py +++ b/elastica/systems/protocol.py @@ -18,11 +18,11 @@ class SystemProtocol(StaticSystemProtocol, Protocol): Protocol for all dynamic elastica system """ - def compute_internal_forces_and_torques(self, time: np.floating) -> None: ... + def compute_internal_forces_and_torques(self, time: np.float64) -> None: ... - def update_accelerations(self, time: np.floating) -> None: ... + def update_accelerations(self, time: np.float64) -> None: ... - def zeroed_out_external_forces_and_torques(self, time: np.floating) -> None: ... + def zeroed_out_external_forces_and_torques(self, time: np.float64) -> None: ... class SlenderBodyGeometryProtocol(Protocol): @@ -32,19 +32,19 @@ def n_nodes(self) -> int: ... @property def n_elems(self) -> int: ... - position_collection: NDArray[np.floating] - velocity_collection: NDArray[np.floating] - acceleration_collection: NDArray[np.floating] + position_collection: NDArray[np.float64] + velocity_collection: NDArray[np.float64] + acceleration_collection: NDArray[np.float64] - omega_collection: NDArray[np.floating] - alpha_collection: NDArray[np.floating] - director_collection: NDArray[np.floating] + omega_collection: NDArray[np.float64] + alpha_collection: NDArray[np.float64] + director_collection: NDArray[np.float64] - external_forces: NDArray[np.floating] - external_torques: NDArray[np.floating] + external_forces: NDArray[np.float64] + external_torques: NDArray[np.float64] - internal_forces: NDArray[np.floating] - internal_torques: NDArray[np.floating] + internal_forces: NDArray[np.float64] + internal_torques: NDArray[np.float64] class SymplecticSystemProtocol(SystemProtocol, SlenderBodyGeometryProtocol, Protocol): @@ -52,8 +52,8 @@ class SymplecticSystemProtocol(SystemProtocol, SlenderBodyGeometryProtocol, Prot Protocol for system with symplectic state variables """ - v_w_collection: NDArray[np.floating] - dvdt_dwdt_collection: NDArray[np.floating] + v_w_collection: NDArray[np.float64] + dvdt_dwdt_collection: NDArray[np.float64] @property def kinematic_states(self) -> _KinematicState: ... @@ -62,18 +62,18 @@ def kinematic_states(self) -> _KinematicState: ... def dynamic_states(self) -> _DynamicState: ... def kinematic_rates( - self, time: np.floating, prefac: np.floating - ) -> tuple[NDArray[np.floating], NDArray[np.floating]]: ... + self, time: np.float64, prefac: np.float64 + ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: ... def dynamic_rates( - self, time: np.floating, prefac: np.floating - ) -> NDArray[np.floating]: ... + self, time: np.float64, prefac: np.float64 + ) -> NDArray[np.float64]: ... class ExplicitSystemProtocol(SystemProtocol, SlenderBodyGeometryProtocol, Protocol): # TODO: Temporarily made to handle explicit stepper. # Need to be refactored as the explicit stepper is further developed. - def __call__(self, time: np.floating, dt: np.floating) -> np.floating: ... + def __call__(self, time: np.float64, dt: np.float64) -> np.float64: ... @property def state(self) -> StateType: ... @state.setter diff --git a/elastica/timestepper/__init__.py b/elastica/timestepper/__init__.py index d2d095365..25f883411 100644 --- a/elastica/timestepper/__init__.py +++ b/elastica/timestepper/__init__.py @@ -19,7 +19,7 @@ def extend_stepper_interface( stepper: StepperProtocol, system_collection: SystemCollectionType ) -> Tuple[ Callable[ - [StepperProtocol, SystemCollectionType, np.floating, np.floating], np.floating + [StepperProtocol, SystemCollectionType, np.float64, np.float64], np.float64 ], SteppersOperatorsType, ]: diff --git a/elastica/timestepper/explicit_steppers.py b/elastica/timestepper/explicit_steppers.py index 5a5877f67..5fb1531f7 100644 --- a/elastica/timestepper/explicit_steppers.py +++ b/elastica/timestepper/explicit_steppers.py @@ -113,9 +113,9 @@ def n_stages(self: ExplicitStepperProtocol) -> int: def step( self: ExplicitStepperProtocol, SystemCollection: SystemCollectionType, - time: np.floating, - dt: np.floating, - ) -> np.floating: + time: np.float64, + dt: np.float64, + ) -> np.float64: if isinstance( self, EulerForward ): # TODO: Cleanup - use depedency injection instead @@ -135,9 +135,9 @@ def do_step( steps_and_prefactors: SteppersOperatorsType, SystemCollection: SystemCollectionType, MemoryCollection: Any, # TODO - time: np.floating, - dt: np.floating, - ) -> np.floating: + time: np.float64, + dt: np.float64, + ) -> np.float64: for stage, update in steps_and_prefactors: SystemCollection.synchronize(time) for system, memory in zip(SystemCollection[:-1], MemoryCollection[:-1]): @@ -152,9 +152,9 @@ def step_single_instance( self: ExplicitStepperProtocol, System: SystemType, Memory: MemoryProtocol, - time: np.floating, - dt: np.floating, - ) -> np.floating: + time: np.float64, + dt: np.float64, + ) -> np.float64: for stage, update in self.steps_and_prefactors: stage(System, Memory, time, dt) time = update(System, Memory, time, dt) @@ -176,8 +176,8 @@ def _first_stage( self, System: ExplicitSystemProtocol, Memory: EulerForwardMemory, - time: np.floating, - dt: np.floating, + time: np.float64, + dt: np.float64, ) -> None: pass @@ -185,9 +185,9 @@ def _first_update( self, System: ExplicitSystemProtocol, Memory: EulerForwardMemory, - time: np.floating, - dt: np.floating, - ) -> np.floating: + time: np.float64, + dt: np.float64, + ) -> np.float64: System.state += dt * System(time, dt) # type: ignore[arg-type] return time + dt @@ -221,8 +221,8 @@ def _first_stage( self, System: ExplicitSystemProtocol, Memory: RungeKutta4Memory, - time: np.floating, - dt: np.floating, + time: np.float64, + dt: np.float64, ) -> None: Memory.initial_state = copy(System.state) Memory.k_1 = dt * System(time, dt) # type: ignore[operator, assignment] @@ -231,9 +231,9 @@ def _first_update( self, System: ExplicitSystemProtocol, Memory: RungeKutta4Memory, - time: np.floating, - dt: np.floating, - ) -> np.floating: + time: np.float64, + dt: np.float64, + ) -> np.float64: # prepare for next stage System.state = Memory.initial_state + 0.5 * Memory.k_1 # type: ignore[operator] return time + 0.5 * dt @@ -242,8 +242,8 @@ def _second_stage( self, System: ExplicitSystemProtocol, Memory: RungeKutta4Memory, - time: np.floating, - dt: np.floating, + time: np.float64, + dt: np.float64, ) -> None: Memory.k_2 = dt * System(time, dt) # type: ignore[operator, assignment] @@ -251,9 +251,9 @@ def _second_update( self, System: ExplicitSystemProtocol, Memory: RungeKutta4Memory, - time: np.floating, - dt: np.floating, - ) -> np.floating: + time: np.float64, + dt: np.float64, + ) -> np.float64: # prepare for next stage System.state = Memory.initial_state + 0.5 * Memory.k_2 # type: ignore[operator] return time @@ -262,8 +262,8 @@ def _third_stage( self, System: ExplicitSystemProtocol, Memory: RungeKutta4Memory, - time: np.floating, - dt: np.floating, + time: np.float64, + dt: np.float64, ) -> None: Memory.k_3 = dt * System(time, dt) # type: ignore[operator, assignment] @@ -271,9 +271,9 @@ def _third_update( self, System: ExplicitSystemProtocol, Memory: RungeKutta4Memory, - time: np.floating, - dt: np.floating, - ) -> np.floating: + time: np.float64, + dt: np.float64, + ) -> np.float64: # prepare for next stage System.state = Memory.initial_state + Memory.k_3 # type: ignore[operator] return time + 0.5 * dt @@ -282,8 +282,8 @@ def _fourth_stage( self, System: ExplicitSystemProtocol, Memory: RungeKutta4Memory, - time: np.floating, - dt: np.floating, + time: np.float64, + dt: np.float64, ) -> None: Memory.k_4 = dt * System(time, dt) # type: ignore[operator, assignment] @@ -291,9 +291,9 @@ def _fourth_update( self, System: ExplicitSystemProtocol, Memory: RungeKutta4Memory, - time: np.floating, - dt: np.floating, - ) -> np.floating: + time: np.float64, + dt: np.float64, + ) -> np.float64: # prepare for next stage System.state = ( Memory.initial_state diff --git a/elastica/timestepper/protocol.py b/elastica/timestepper/protocol.py index 6eff40dd9..1a64c7254 100644 --- a/elastica/timestepper/protocol.py +++ b/elastica/timestepper/protocol.py @@ -25,12 +25,12 @@ def n_stages(self) -> int: ... def step_methods(self) -> SteppersOperatorsType: ... def step( - self, SystemCollection: SystemCollectionType, time: np.floating, dt: np.floating - ) -> np.floating: ... + self, SystemCollection: SystemCollectionType, time: np.float64, dt: np.float64 + ) -> np.float64: ... def step_single_instance( - self, SystemCollection: SystemType, time: np.floating, dt: np.floating - ) -> np.floating: ... + self, SystemCollection: SystemType, time: np.float64, dt: np.float64 + ) -> np.float64: ... class SymplecticStepperProtocol(StepperProtocol, Protocol): diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index 5d7c0ca37..ae6ace3f7 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -69,9 +69,9 @@ def n_stages(self: SymplecticStepperProtocol) -> int: def step( self: SymplecticStepperProtocol, SystemCollection: SystemCollectionType, - time: np.floating, - dt: np.floating, - ) -> np.floating: + time: np.float64, + dt: np.float64, + ) -> np.float64: return SymplecticStepperMixin.do_step(self, self.steps_and_prefactors, SystemCollection, time, dt) # type: ignore[attr-defined] # TODO: Merge with .step method in the future. @@ -81,9 +81,9 @@ def do_step( TimeStepper: SymplecticStepperProtocol, steps_and_prefactors: SteppersOperatorsType, SystemCollection: SystemCollectionType, - time: np.floating, - dt: np.floating, - ) -> np.floating: + time: np.float64, + dt: np.float64, + ) -> np.float64: """ Function for doing symplectic stepper over the user defined rods (system). @@ -138,9 +138,9 @@ def do_step( def step_single_instance( self: SymplecticStepperProtocol, System: SymplecticSystemProtocol, - time: np.floating, - dt: np.floating, - ) -> np.floating: + time: np.float64, + dt: np.float64, + ) -> np.float64: for kin_prefactor, kin_step, dyn_step in self.steps_and_prefactors[:-1]: kin_step(System, time, dt) @@ -175,11 +175,11 @@ def get_prefactors(self) -> list[OperatorType]: self._first_prefactor, ] - def _first_prefactor(self, dt: np.floating) -> np.floating: + def _first_prefactor(self, dt: np.float64) -> np.float64: return 0.5 * dt def _first_kinematic_step( - self, System: SymplecticSystemProtocol, time: np.floating, dt: np.floating + self, System: SymplecticSystemProtocol, time: np.float64, dt: np.float64 ) -> None: prefac = self._first_prefactor(dt) overload_operator_kinematic_numba( @@ -192,7 +192,7 @@ def _first_kinematic_step( ) def _first_dynamic_step( - self, System: SymplecticSystemProtocol, time: np.floating, dt: np.floating + self, System: SymplecticSystemProtocol, time: np.float64, dt: np.float64 ) -> None: overload_operator_dynamic_numba( System.dynamic_states.rate_collection, @@ -235,11 +235,11 @@ def get_prefactors(self) -> list[OperatorType]: self._first_kinematic_prefactor, ] - def _first_kinematic_prefactor(self, dt: np.floating) -> np.floating: + def _first_kinematic_prefactor(self, dt: np.float64) -> np.float64: return self.ξ * dt def _first_kinematic_step( - self, System: SymplecticSystemProtocol, time: np.floating, dt: np.floating + self, System: SymplecticSystemProtocol, time: np.float64, dt: np.float64 ) -> None: prefac = self._first_kinematic_prefactor(dt) overload_operator_kinematic_numba( @@ -253,7 +253,7 @@ def _first_kinematic_step( # System.kinematic_states += prefac * System.kinematic_rates(time, prefac) def _first_dynamic_step( - self, System: SymplecticSystemProtocol, time: np.floating, dt: np.floating + self, System: SymplecticSystemProtocol, time: np.float64, dt: np.float64 ) -> None: prefac = self.lambda_dash_coeff * dt overload_operator_dynamic_numba( @@ -262,11 +262,11 @@ def _first_dynamic_step( ) # System.dynamic_states += prefac * System.dynamic_rates(time, prefac) - def _second_kinematic_prefactor(self, dt: np.floating) -> np.floating: + def _second_kinematic_prefactor(self, dt: np.float64) -> np.float64: return self.χ * dt def _second_kinematic_step( - self, System: SymplecticSystemProtocol, time: np.floating, dt: np.floating + self, System: SymplecticSystemProtocol, time: np.float64, dt: np.float64 ) -> None: prefac = self._second_kinematic_prefactor(dt) overload_operator_kinematic_numba( @@ -280,7 +280,7 @@ def _second_kinematic_step( # System.kinematic_states += prefac * System.kinematic_rates(time, prefac) def _second_dynamic_step( - self, System: SymplecticSystemProtocol, time: np.floating, dt: np.floating + self, System: SymplecticSystemProtocol, time: np.float64, dt: np.float64 ) -> None: prefac = self.λ * dt overload_operator_dynamic_numba( @@ -289,11 +289,11 @@ def _second_dynamic_step( ) # System.dynamic_states += prefac * System.dynamic_rates(time, prefac) - def _third_kinematic_prefactor(self, dt: np.floating) -> np.floating: + def _third_kinematic_prefactor(self, dt: np.float64) -> np.float64: return self.xi_chi_dash_coeff * dt def _third_kinematic_step( - self, System: SymplecticSystemProtocol, time: np.floating, dt: np.floating + self, System: SymplecticSystemProtocol, time: np.float64, dt: np.float64 ) -> None: prefac = self._third_kinematic_prefactor(dt) # Need to fill in diff --git a/elastica/transformations.py b/elastica/transformations.py index c8e977211..3a2e10b5a 100644 --- a/elastica/transformations.py +++ b/elastica/transformations.py @@ -16,8 +16,8 @@ def format_vector_shape( - vector_collection: NDArray[np.floating], -) -> NDArray[np.floating]: + vector_collection: NDArray[np.float64], +) -> NDArray[np.float64]: """ Function for formatting vector shapes into correct format @@ -63,8 +63,8 @@ def format_vector_shape( def format_matrix_shape( - matrix_collection: NDArray[np.floating], -) -> NDArray[np.floating]: + matrix_collection: NDArray[np.float64], +) -> NDArray[np.float64]: """ Formats input matrix into correct format @@ -141,14 +141,14 @@ def assert_proper_square(num1: int) -> int: return matrix_collection -def skew_symmetrize(vector: NDArray[np.floating]) -> NDArray[np.floating]: +def skew_symmetrize(vector: NDArray[np.float64]) -> NDArray[np.float64]: vector = format_vector_shape(vector) return _skew_symmetrize(vector) def inv_skew_symmetrize( - matrix_collection: NDArray[np.floating], -) -> NDArray[np.floating]: + matrix_collection: NDArray[np.float64], +) -> NDArray[np.float64]: """ Safe wrapper around inv_skew_symmetrize that does checking and formatting on type of matrix_collection using format_matrix_shape @@ -175,8 +175,8 @@ def inv_skew_symmetrize( def rotate( - matrix: NDArray[np.floating], scale: np.floating, axis: NDArray[np.floating] -) -> NDArray[np.floating]: + matrix: NDArray[np.float64], scale: np.float64, axis: NDArray[np.float64] +) -> NDArray[np.float64]: """ This function takes single or multiple frames as matrix. Then rotates these frames around a single axis for all frames, or can rotate each frame around its own diff --git a/elastica/utils.py b/elastica/utils.py index 429ae63ab..7f851ed05 100644 --- a/elastica/utils.py +++ b/elastica/utils.py @@ -179,7 +179,7 @@ def extend_instance(obj: Any, cls: Any) -> None: def _bspline( # type: ignore[no-any-unimported] - t_coeff: NDArray, l_centerline: np.floating = np.float64(1.0) + t_coeff: NDArray, l_centerline: np.float64 = np.float64(1.0) ) -> tuple[BSpline, NDArray, NDArray]: """Generates a bspline object that plots the spline interpolant for any vector x. Optionally takes in a centerline length, set to 1.0 by diff --git a/examples/MuscularSnake/muscular_snake.py b/examples/MuscularSnake/muscular_snake.py index 2d224c63f..61d68f304 100644 --- a/examples/MuscularSnake/muscular_snake.py +++ b/examples/MuscularSnake/muscular_snake.py @@ -151,8 +151,8 @@ class MuscularSnakeSimulator( muscle_glue_connection_index.append( np.hstack( ( - np.arange(0, 4 * 3, 1, dtype=np.int64), - np.arange(9 * 3, n_elem_muscle_group_one_to_three, 1, dtype=np.int64), + np.arange(0, 4 * 3, 1, dtype=np.int32), + np.arange(9 * 3, n_elem_muscle_group_one_to_three, 1, dtype=np.int32), ) ) ) @@ -213,11 +213,11 @@ class MuscularSnakeSimulator( muscle_rod_list.append(muscle_rod) muscle_end_connection_index.append(index + n_elem_muscle_group_two_to_four) muscle_glue_connection_index.append( - # np.array([0,1, 2, 3, 9, 10 ], dtype=np.int) + # np.array([0,1, 2, 3, 9, 10 ], dtype=np.int32) np.hstack( ( - np.arange(0, 4 * 3, 1, dtype=np.int64), - np.arange(9 * 3, n_elem_muscle_group_two_to_four, 1, dtype=np.int64), + np.arange(0, 4 * 3, 1, dtype=np.int32), + np.arange(9 * 3, n_elem_muscle_group_two_to_four, 1, dtype=np.int32), ) ) ) diff --git a/examples/RigidBodyCases/rigid_cylinder_rotational_motion_case.py b/examples/RigidBodyCases/rigid_cylinder_rotational_motion_case.py index f33b6ccd0..642719166 100644 --- a/examples/RigidBodyCases/rigid_cylinder_rotational_motion_case.py +++ b/examples/RigidBodyCases/rigid_cylinder_rotational_motion_case.py @@ -52,7 +52,7 @@ def __init__(self, torque, direction=np.array([0.0, 0.0, 0.0])): super(PointCoupleToCenter, self).__init__() self.torque = (torque * direction).reshape(3, 1) - def apply_forces(self, system, time: np.floating = np.float64(0.0)): + def apply_forces(self, system, time: np.float64 = np.float64(0.0)): system.external_torques += np.einsum( "ijk, jk->ik", system.director_collection, self.torque ) diff --git a/examples/RigidBodyCases/rigid_cylinder_translational_motion_case.py b/examples/RigidBodyCases/rigid_cylinder_translational_motion_case.py index 0e21bc0ba..ce78a84af 100644 --- a/examples/RigidBodyCases/rigid_cylinder_translational_motion_case.py +++ b/examples/RigidBodyCases/rigid_cylinder_translational_motion_case.py @@ -52,7 +52,7 @@ def __init__(self, force, direction=np.array([0.0, 0.0, 0.0])): super(PointForceToCenter, self).__init__() self.force = (force * direction).reshape(3, 1) - def apply_forces(self, system, time: np.floating = np.float64(0.0)): + def apply_forces(self, system, time: np.float64 = np.float64(0.0)): system.external_forces += self.force # Add point force on the rod diff --git a/examples/RigidBodyCases/rigid_sphere_rotational_motion_case.py b/examples/RigidBodyCases/rigid_sphere_rotational_motion_case.py index 2f377209a..600d93d87 100644 --- a/examples/RigidBodyCases/rigid_sphere_rotational_motion_case.py +++ b/examples/RigidBodyCases/rigid_sphere_rotational_motion_case.py @@ -43,7 +43,7 @@ def __init__(self, torque, direction=np.array([0.0, 0.0, 0.0])): super(PointCoupleToCenter, self).__init__() self.torque = (torque * direction).reshape(3, 1) - def apply_forces(self, system, time: np.floating = np.float64(0.0)): + def apply_forces(self, system, time: np.float64 = np.float64(0.0)): system.external_torques += np.einsum( "ijk, jk->ik", system.director_collection, self.torque ) diff --git a/examples/RigidBodyCases/rigid_sphere_translational_motion_case.py b/examples/RigidBodyCases/rigid_sphere_translational_motion_case.py index 3ececd33e..91a190150 100644 --- a/examples/RigidBodyCases/rigid_sphere_translational_motion_case.py +++ b/examples/RigidBodyCases/rigid_sphere_translational_motion_case.py @@ -44,7 +44,7 @@ def __init__(self, force, direction=np.array([0.0, 0.0, 0.0])): super(PointForceToCenter, self).__init__() self.force = (force * direction).reshape(3, 1) - def apply_forces(self, system, time: np.floating = np.float64(0.0)): + def apply_forces(self, system, time: np.float64 = np.float64(0.0)): system.external_forces += self.force # Add point force on the rod diff --git a/tests/test_math/test_governing_equations.py b/tests/test_math/test_governing_equations.py index eca87eefb..92c223325 100644 --- a/tests/test_math/test_governing_equations.py +++ b/tests/test_math/test_governing_equations.py @@ -49,8 +49,8 @@ def constructor(n_elem): ) # Ghost needed for Cosserat rod functions adapted for block structure. - rod.ghost_elems_idx = np.empty((0), dtype=int) - rod.ghost_voronoi_idx = np.empty((0), dtype=int) + rod.ghost_elems_idx = np.empty((0), dtype=np.int32) + rod.ghost_voronoi_idx = np.empty((0), dtype=np.int32) return cls, rod @@ -372,7 +372,7 @@ def test_case_compute_internal_forces(self, n_elem, dilatation): test_rod.shear_matrix, test_rod.internal_stress, test_rod.internal_forces, - ghost_elems_idx=np.empty((0), dtype=int), + ghost_elems_idx=np.empty((0), dtype=np.int32), ) assert_allclose( @@ -725,7 +725,7 @@ def test_case_lagrange_transport_unsteady_dilatation(self): test_rod.dilatation, test_rod.dilatation_rate, test_rod.internal_torques, - ghost_voronoi_idx=np.empty((0), dtype=int), + ghost_voronoi_idx=np.empty((0), dtype=np.int32), ) # computed internal torques has to be zero. Internal torques created by Lagrangian diff --git a/tests/test_memory_block_validity.py b/tests/test_memory_block_validity.py index 12fbd2346..0c2649543 100644 --- a/tests/test_memory_block_validity.py +++ b/tests/test_memory_block_validity.py @@ -749,10 +749,10 @@ def test_periodic_boundary_one_ring_rod(): block_structure = MemoryBlockCosseratRod([ring_rod], [0]) correct_periodic_boundary_node_idx = np.array( - [[0, 6, 7], [5, 1, 2]], dtype=np.int64 + [[0, 6, 7], [5, 1, 2]], dtype=np.int32 ) - correct_periodic_boundary_elem_idx = np.array([[0, 6], [5, 1]], dtype=np.int64) - correct_periodic_boundary_voronoi_idx = np.array([[0], [5]], dtype=np.int64) + correct_periodic_boundary_elem_idx = np.array([[0, 6], [5, 1]], dtype=np.int32) + correct_periodic_boundary_voronoi_idx = np.array([[0], [5]], dtype=np.int32) assert_allclose( correct_periodic_boundary_node_idx, @@ -770,14 +770,14 @@ def test_periodic_boundary_one_ring_rod(): atol=Tolerance.atol(), ) - correct_start_node_idx = np.array([1], dtype=np.int64) - correct_end_node_idx = np.array([6], dtype=np.int64) + correct_start_node_idx = np.array([1], dtype=np.int32) + correct_end_node_idx = np.array([6], dtype=np.int32) - correct_start_elem_idx = np.array([1], dtype=np.int64) - correct_end_elem_idx = np.array([6], dtype=np.int64) + correct_start_elem_idx = np.array([1], dtype=np.int32) + correct_end_elem_idx = np.array([6], dtype=np.int32) - correct_start_voronoi_idx = np.array([1], dtype=np.int64) - correct_end_voronoi_idx = np.array([6], dtype=np.int64) + correct_start_voronoi_idx = np.array([1], dtype=np.int32) + correct_end_voronoi_idx = np.array([6], dtype=np.int32) assert_allclose( correct_start_node_idx, @@ -826,12 +826,12 @@ def test_periodic_boundary_two_ring_rod(): block_structure = MemoryBlockCosseratRod([ring_rod_1, ring_rod_2], [0, 1]) correct_periodic_boundary_node_idx = np.array( - [[0, 6, 7, 9, 15, 16], [5, 1, 2, 14, 10, 11]], dtype=np.int64 + [[0, 6, 7, 9, 15, 16], [5, 1, 2, 14, 10, 11]], dtype=np.int32 ) correct_periodic_boundary_elem_idx = np.array( - [[0, 6, 9, 15], [5, 1, 14, 10]], dtype=np.int64 + [[0, 6, 9, 15], [5, 1, 14, 10]], dtype=np.int32 ) - correct_periodic_boundary_voronoi_idx = np.array([[0, 9], [5, 14]], dtype=np.int64) + correct_periodic_boundary_voronoi_idx = np.array([[0, 9], [5, 14]], dtype=np.int32) assert_allclose( correct_periodic_boundary_node_idx, @@ -849,14 +849,14 @@ def test_periodic_boundary_two_ring_rod(): atol=Tolerance.atol(), ) - correct_start_node_idx = np.array([1, 10], dtype=np.int64) - correct_end_node_idx = np.array([6, 15], dtype=np.int64) + correct_start_node_idx = np.array([1, 10], dtype=np.int32) + correct_end_node_idx = np.array([6, 15], dtype=np.int32) - correct_start_elem_idx = np.array([1, 10], dtype=np.int64) - correct_end_elem_idx = np.array([6, 15], dtype=np.int64) + correct_start_elem_idx = np.array([1, 10], dtype=np.int32) + correct_end_elem_idx = np.array([6, 15], dtype=np.int32) - correct_start_voronoi_idx = np.array([1, 10], dtype=np.int64) - correct_end_voronoi_idx = np.array([6, 15], dtype=np.int64) + correct_start_voronoi_idx = np.array([1, 10], dtype=np.int32) + correct_end_voronoi_idx = np.array([6, 15], dtype=np.int32) assert_allclose( correct_start_node_idx, @@ -906,12 +906,12 @@ def test_periodic_boundary_two_ring_rod_different_nelems(): block_structure = MemoryBlockCosseratRod([ring_rod_1, ring_rod_2], [0, 1]) correct_periodic_boundary_node_idx = np.array( - [[0, 6, 7, 9, 13, 14], [5, 1, 2, 12, 10, 11]], dtype=np.int64 + [[0, 6, 7, 9, 13, 14], [5, 1, 2, 12, 10, 11]], dtype=np.int32 ) correct_periodic_boundary_elem_idx = np.array( - [[0, 6, 9, 13], [5, 1, 12, 10]], dtype=np.int64 + [[0, 6, 9, 13], [5, 1, 12, 10]], dtype=np.int32 ) - correct_periodic_boundary_voronoi_idx = np.array([[0, 9], [5, 12]], dtype=np.int64) + correct_periodic_boundary_voronoi_idx = np.array([[0, 9], [5, 12]], dtype=np.int32) assert_allclose( correct_periodic_boundary_node_idx, @@ -929,14 +929,14 @@ def test_periodic_boundary_two_ring_rod_different_nelems(): atol=Tolerance.atol(), ) - correct_start_node_idx = np.array([1, 10], dtype=np.int64) - correct_end_node_idx = np.array([6, 13], dtype=np.int64) + correct_start_node_idx = np.array([1, 10], dtype=np.int32) + correct_end_node_idx = np.array([6, 13], dtype=np.int32) - correct_start_elem_idx = np.array([1, 10], dtype=np.int64) - correct_end_elem_idx = np.array([6, 13], dtype=np.int64) + correct_start_elem_idx = np.array([1, 10], dtype=np.int32) + correct_end_elem_idx = np.array([6, 13], dtype=np.int32) - correct_start_voronoi_idx = np.array([1, 10], dtype=np.int64) - correct_end_voronoi_idx = np.array([6, 13], dtype=np.int64) + correct_start_voronoi_idx = np.array([1, 10], dtype=np.int32) + correct_end_voronoi_idx = np.array([6, 13], dtype=np.int32) assert_allclose( correct_start_node_idx, @@ -985,10 +985,10 @@ def test_periodic_boundary_one_ring_one_straight_rod(): block_structure = MemoryBlockCosseratRod([ring_rod, straight_rod], [0, 1]) correct_periodic_boundary_node_idx = np.array( - [[7, 13, 14], [12, 8, 9]], dtype=np.int64 + [[7, 13, 14], [12, 8, 9]], dtype=np.int32 ) - correct_periodic_boundary_elem_idx = np.array([[7, 13], [12, 8]], dtype=np.int64) - correct_periodic_boundary_voronoi_idx = np.array([[7], [12]], dtype=np.int64) + correct_periodic_boundary_elem_idx = np.array([[7, 13], [12, 8]], dtype=np.int32) + correct_periodic_boundary_voronoi_idx = np.array([[7], [12]], dtype=np.int32) assert_allclose( correct_periodic_boundary_node_idx, @@ -1006,14 +1006,14 @@ def test_periodic_boundary_one_ring_one_straight_rod(): atol=Tolerance.atol(), ) - correct_start_node_idx = np.array([0, 8], dtype=np.int64) - correct_end_node_idx = np.array([6, 13], dtype=np.int64) + correct_start_node_idx = np.array([0, 8], dtype=np.int32) + correct_end_node_idx = np.array([6, 13], dtype=np.int32) - correct_start_elem_idx = np.array([0, 8], dtype=np.int64) - correct_end_elem_idx = np.array([5, 13], dtype=np.int64) + correct_start_elem_idx = np.array([0, 8], dtype=np.int32) + correct_end_elem_idx = np.array([5, 13], dtype=np.int32) - correct_start_voronoi_idx = np.array([0, 8], dtype=np.int64) - correct_end_voronoi_idx = np.array([4, 13], dtype=np.int64) + correct_start_voronoi_idx = np.array([0, 8], dtype=np.int32) + correct_end_voronoi_idx = np.array([4, 13], dtype=np.int32) assert_allclose( correct_start_node_idx, diff --git a/tests/test_modules/test_memory_block_utils.py b/tests/test_modules/test_memory_block_utils.py index 580addabd..663a3aa27 100644 --- a/tests/test_modules/test_memory_block_utils.py +++ b/tests/test_modules/test_memory_block_utils.py @@ -13,48 +13,48 @@ "n_elems_in_rods, outputs", [ ( - np.array([5], dtype=np.int64), + np.array([5], dtype=np.int32), [ np.int64(5), - np.array([], dtype=np.int64), - np.array([], dtype=np.int64), - np.array([], dtype=np.int64), + np.array([], dtype=np.int32), + np.array([], dtype=np.int32), + np.array([], dtype=np.int32), ], ), ( - np.array([5, 5], dtype=np.int64), + np.array([5, 5], dtype=np.int32), [ np.int64(12), - np.array([6], dtype=np.int64), - np.array([5, 6], dtype=np.int64), - np.array([4, 5, 6], dtype=np.int64), + np.array([6], dtype=np.int32), + np.array([5, 6], dtype=np.int32), + np.array([4, 5, 6], dtype=np.int32), ], ), ( - np.array([1, 1, 1], dtype=np.int64), + np.array([1, 1, 1], dtype=np.int32), [ np.int64(7), - np.array([2, 5], dtype=np.int64), - np.array([1, 2, 4, 5], dtype=np.int64), - np.array([0, 1, 2, 3, 4, 5], dtype=np.int64), + np.array([2, 5], dtype=np.int32), + np.array([1, 2, 4, 5], dtype=np.int32), + np.array([0, 1, 2, 3, 4, 5], dtype=np.int32), ], ), ( - np.array([1, 2, 3], dtype=np.int64), + np.array([1, 2, 3], dtype=np.int32), [ np.int64(10), - np.array([2, 6], dtype=np.int64), - np.array([1, 2, 5, 6], dtype=np.int64), - np.array([0, 1, 2, 4, 5, 6], dtype=np.int64), + np.array([2, 6], dtype=np.int32), + np.array([1, 2, 5, 6], dtype=np.int32), + np.array([0, 1, 2, 4, 5, 6], dtype=np.int32), ], ), ( - np.array([10, 10, 5, 5], dtype=np.int64), + np.array([10, 10, 5, 5], dtype=np.int32), [ np.int64(36), - np.array([11, 23, 30], dtype=np.int64), - np.array([10, 11, 22, 23, 29, 30], dtype=np.int64), - np.array([9, 10, 11, 21, 22, 23, 28, 29, 30], dtype=np.int64), + np.array([11, 23, 30], dtype=np.int32), + np.array([10, 11, 22, 23, 29, 30], dtype=np.int32), + np.array([9, 10, 11, 21, 22, 23, 28, 29, 30], dtype=np.int32), ], ), ], @@ -76,10 +76,10 @@ def test_make_block_memory_metadata(n_elems_in_rods, outputs): @pytest.mark.parametrize( "n_elems_in_ring_rods", [ - np.array([10], dtype=np.int64), - np.array([2, 4], dtype=np.int64), - np.array([4, 5, 7], dtype=np.int64), - np.array([10, 10, 10, 10], dtype=np.int64), + np.array([10], dtype=np.int32), + np.array([2, 4], dtype=np.int32), + np.array([4, 5, 7], dtype=np.int32), + np.array([10, 10, 10, 10], dtype=np.int32), ], ) def test_make_block_memory_periodic_boundary_metadata(n_elems_in_ring_rods): @@ -92,9 +92,9 @@ def test_make_block_memory_periodic_boundary_metadata(n_elems_in_ring_rods): n_ring_rods = len(n_elems_in_ring_rods) expected_n_elem = n_elems_in_ring_rods + 2 - expected_node_idx = np.empty((2, 3 * n_ring_rods), dtype=np.int64) - expected_element_idx = np.empty((2, 2 * n_ring_rods), dtype=np.int64) - expected_voronoi_idx = np.empty((2, n_ring_rods), dtype=np.int64) + expected_node_idx = np.empty((2, 3 * n_ring_rods), dtype=np.int32) + expected_element_idx = np.empty((2, 2 * n_ring_rods), dtype=np.int32) + expected_voronoi_idx = np.empty((2, n_ring_rods), dtype=np.int32) accumulation = np.hstack((0, np.cumsum(n_elems_in_ring_rods[:-1] + 4))) diff --git a/tests/test_synchronize_periodic_boundary.py b/tests/test_synchronize_periodic_boundary.py index 08a1787c1..b20d776dd 100644 --- a/tests/test_synchronize_periodic_boundary.py +++ b/tests/test_synchronize_periodic_boundary.py @@ -28,7 +28,7 @@ def test_synchronize_periodic_boundary_vector(n_elems): input_vector = np.random.random((3, n_elems + 3)) - periodic_idx = np.zeros((2, 3), dtype=np.int64) + periodic_idx = np.zeros((2, 3), dtype=np.int32) periodic_idx[0, 0] = 0 periodic_idx[0, 1] = -2 periodic_idx[0, 2] = -1 @@ -63,7 +63,7 @@ def test_synchronize_periodic_boundary_matrix(n_elems): input_matrix = np.random.random((3, 3, n_elems + 3)) - periodic_idx = np.zeros((2, 3), dtype=np.int64) + periodic_idx = np.zeros((2, 3), dtype=np.int32) periodic_idx[0, 0] = 0 periodic_idx[0, 1] = -2 periodic_idx[0, 2] = -1 @@ -98,7 +98,7 @@ def test_synchronize_periodic_boundary_scalar(n_elems): input_matrix = np.random.random((n_elems + 3)) - periodic_idx = np.zeros((2, 3), dtype=np.int64) + periodic_idx = np.zeros((2, 3), dtype=np.int32) periodic_idx[0, 0] = 0 periodic_idx[0, 1] = -2 periodic_idx[0, 2] = -1 From 76f642c8f6e73ae887c94b3adb6a277d814c5e83 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Tue, 25 Jun 2024 22:26:38 -0500 Subject: [PATCH 128/134] test: fix few typings that has been ignored --- elastica/external_forces.py | 19 +++++++++++++------ elastica/memory_block/memory_block_rod.py | 6 +++--- elastica/timestepper/symplectic_steppers.py | 6 ++++-- 3 files changed, 20 insertions(+), 11 deletions(-) diff --git a/elastica/external_forces.py b/elastica/external_forces.py index d7f5185ec..ad15810ee 100644 --- a/elastica/external_forces.py +++ b/elastica/external_forces.py @@ -77,7 +77,10 @@ class GravityForces(NoForces): """ def __init__( - self, acc_gravity: NDArray[np.float64] = np.array([0.0, -9.80665, 0.0]) + self, + acc_gravity: NDArray[np.float64] = np.array( + [0.0, -9.80665, 0.0] + ), # FIXME: avoid mutable default ) -> None: """ @@ -199,9 +202,9 @@ def compute_end_point_forces( Applied forces are ramped up until ramp up time. """ - factor = min(1.0, time / ramp_up_time) - external_forces[..., 0] += start_force * factor # type: ignore[operator] - external_forces[..., -1] += end_force * factor # type: ignore[operator] + factor = min(np.float64(1.0), time / ramp_up_time) + external_forces[..., 0] += start_force * factor + external_forces[..., -1] += end_force * factor class UniformTorques(NoForces): @@ -218,7 +221,9 @@ class UniformTorques(NoForces): def __init__( self, torque: np.float64, - direction: NDArray[np.float64] = np.array([0.0, 0.0, 0.0]), + direction: NDArray[np.float64] = np.array( + [0.0, 0.0, 0.0] + ), # FIXME: avoid mutable default ) -> None: """ @@ -258,7 +263,9 @@ class UniformForces(NoForces): def __init__( self, force: np.float64, - direction: NDArray[np.float64] = np.array([0.0, 0.0, 0.0]), + direction: NDArray[np.float64] = np.array( + [0.0, 0.0, 0.0] + ), # FIXME: avoid mutable default ) -> None: """ diff --git a/elastica/memory_block/memory_block_rod.py b/elastica/memory_block/memory_block_rod.py index a363f3d3c..a47998bac 100644 --- a/elastica/memory_block/memory_block_rod.py +++ b/elastica/memory_block/memory_block_rod.py @@ -134,13 +134,13 @@ def __init__( # we place first straight rods, then ring rods. # TODO: in future consider a better implementation for packing problem. # +1 is because we want to start from next idx, where periodic boundary starts - self.periodic_boundary_nodes_idx += ( # type: ignore + self.periodic_boundary_nodes_idx += ( self.ghost_nodes_idx[n_straight_rods - 1] + 1 ) - self.periodic_boundary_elems_idx += ( # type: ignore + self.periodic_boundary_elems_idx += ( self.ghost_elems_idx[1::2][n_straight_rods - 1] + 1 ) - self.periodic_boundary_voronoi_idx += ( # type: ignore + self.periodic_boundary_voronoi_idx += ( self.ghost_voronoi_idx[2::3][n_straight_rods - 1] + 1 ) diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index ae6ace3f7..29c3a7e06 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -72,7 +72,9 @@ def step( time: np.float64, dt: np.float64, ) -> np.float64: - return SymplecticStepperMixin.do_step(self, self.steps_and_prefactors, SystemCollection, time, dt) # type: ignore[attr-defined] + return SymplecticStepperMixin.do_step( + self, self.steps_and_prefactors, SystemCollection, time, dt + ) # TODO: Merge with .step method in the future. # DEPRECATED: Use .step instead. @@ -153,7 +155,7 @@ def step_single_instance( last_kin_step = self.steps_and_prefactors[-1][1] last_kin_step(System, time, dt) - return time + last_kin_prefactor(dt) # type: ignore[no-any-return] + return time + last_kin_prefactor(dt) class PositionVerlet(SymplecticStepperMixin): From 9fdc9f057ea5733ecb1e524162d61b1736b5c8bc Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Tue, 25 Jun 2024 22:33:27 -0500 Subject: [PATCH 129/134] add ring-rod typehint check --- elastica/rod/cosserat_rod.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/elastica/rod/cosserat_rod.py b/elastica/rod/cosserat_rod.py index 1027d2266..03309aef9 100644 --- a/elastica/rod/cosserat_rod.py +++ b/elastica/rod/cosserat_rod.py @@ -1140,3 +1140,13 @@ def _zeroed_out_external_forces_and_torques( 1.0, youngs_modulus=1.0, ) + _: CosseratRodProtocol = CosseratRod.ring_rod( # type: ignore[no-redef] + 3, + np.zeros(3), + np.array([0, 1, 0]), + np.array([0, 0, 1]), + 1.0, + 0.1, + 1.0, + youngs_modulus=1.0, + ) From 7039babc5bfa9d3efe4cea46895742bd88d797a6 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Tue, 25 Jun 2024 23:31:11 -0500 Subject: [PATCH 130/134] refactor: clarify iterating `systems` vs `block_systems` --- elastica/modules/base_system.py | 72 +++++++++++---------- elastica/modules/callbacks.py | 4 +- elastica/modules/connections.py | 8 +-- elastica/modules/constraints.py | 14 ++-- elastica/modules/contact.py | 12 ++-- elastica/modules/damping.py | 6 +- elastica/modules/forcing.py | 6 +- elastica/modules/memory_block.py | 3 +- elastica/modules/protocol.py | 11 ++-- elastica/systems/analytical.py | 3 + elastica/timestepper/symplectic_steppers.py | 10 +-- tests/test_modules/test_base_system.py | 22 ++++--- tests/test_modules/test_callbacks.py | 2 +- tests/test_modules/test_connections.py | 22 +++---- tests/test_modules/test_constraints.py | 6 +- tests/test_modules/test_contact.py | 10 +-- tests/test_modules/test_damping.py | 6 +- tests/test_modules/test_forcing.py | 2 +- tests/test_restart.py | 6 +- 19 files changed, 117 insertions(+), 108 deletions(-) diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index 10c111ce6..43666bf66 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -8,8 +8,9 @@ from typing import Type, Generator, Iterable, Any, overload from typing import final from elastica.typing import ( - StaticSystemType, SystemType, + StaticSystemType, + BlockSystemType, SystemIdxType, OperatorType, OperatorCallbackType, @@ -38,8 +39,10 @@ class BaseSystemCollection(MutableSequence): ---------- allowed_sys_types: tuple Tuple of allowed type rod-like objects. Here use a base class for objects, i.e. RodBase. - _systems: list - List of rod-like objects. + systems: Callabke + Returns all system objects. Once finalize, block objects are also included. + blocks: Callable + Returns block objects. Should be called after finalize. Note ---- @@ -68,8 +71,8 @@ def __init__(self) -> None: ) # List of systems to be integrated - self._systems: list[StaticSystemType] = [] - self.__final_systems: list[SystemType] = [] + self.__systems: list[StaticSystemType] = [] + self.__final_blocks: list[BlockSystemType] = [] # Flag Finalize: Finalizing twice will cause an error, # but the error message is very misleading @@ -99,7 +102,7 @@ def _check_type(self, sys_to_be_added: Any) -> bool: return True def __len__(self) -> int: - return len(self._systems) + return len(self.__systems) @overload def __getitem__(self, idx: int, /) -> SystemType: ... @@ -108,22 +111,22 @@ def __getitem__(self, idx: int, /) -> SystemType: ... def __getitem__(self, idx: slice, /) -> list[SystemType]: ... def __getitem__(self, idx, /): # type: ignore - return self._systems[idx] + return self.__systems[idx] def __delitem__(self, idx, /): # type: ignore - del self._systems[idx] + del self.__systems[idx] def __setitem__(self, idx, system, /): # type: ignore self._check_type(system) - self._systems[idx] = system + self.__systems[idx] = system def insert(self, idx, system) -> None: # type: ignore self._check_type(system) - self._systems.insert(idx, system) + self.__systems.insert(idx, system) def __str__(self) -> str: """To be readable""" - return str(self._systems) + return str(self.__systems) @final def extend_allowed_types( @@ -138,38 +141,43 @@ def override_allowed_types( self.allowed_sys_types = allowed_types @final - def _get_sys_idx_if_valid( - self, sys_to_be_added: "SystemType | StaticSystemType" + def get_system_index( + self, system: "SystemType | StaticSystemType" ) -> SystemIdxType: n_systems = len(self) # Total number of systems from mixed-in class sys_idx: SystemIdxType - if isinstance(sys_to_be_added, (int, np.int_)): + if isinstance(system, (int, np.int_)): # 1. If they are indices themselves, check range + # This is only used for testing purposes assert ( - -n_systems <= sys_to_be_added < n_systems - ), "Rod index {} exceeds number of registered rodtems".format( - sys_to_be_added - ) - sys_idx = int(sys_to_be_added) - elif self._check_type(sys_to_be_added): - # 2. If they are rod objects (most likely), lookup indices + -n_systems <= system < n_systems + ), "System index {} exceeds number of registered rodtems".format(system) + sys_idx = int(system) + elif self._check_type(system): + # 2. If they are system object (most likely), lookup indices # index might have some problems : https://stackoverflow.com/a/176921 try: - sys_idx = self._systems.index(sys_to_be_added) + sys_idx = self.__systems.index(system) except ValueError: raise ValueError( - "Rod {} was not found, did you append it to the system?".format( - sys_to_be_added + "System {} was not found, did you append it to the system?".format( + system ) ) return sys_idx @final - def systems(self) -> Generator[SystemType, None, None]: + def systems(self) -> Generator[StaticSystemType, None, None]: + # assert self._finalize_flag, "The simulator is not finalized." + for system in self.__systems: + yield system + + @final + def block_systems(self) -> Generator[BlockSystemType, None, None]: # assert self._finalize_flag, "The simulator is not finalized." - for block in self.__final_systems: + for block in self.__final_blocks: yield block @final @@ -184,15 +192,11 @@ def finalize(self) -> None: assert not self._finalize_flag, "The finalize cannot be called twice." self._finalize_flag = True - # construct memory block - self.__final_systems = construct_memory_block_structures(self._systems) - self._systems.extend( - self.__final_systems - ) # FIXME: We need this to make ring-rod working. + # Construct memory block + self.__final_blocks = construct_memory_block_structures(self.__systems) + # FIXME: We need this to make ring-rod working. # But probably need to be refactored - # TODO: try to remove the _systems list for memory optimization - # self._systems.clear() - # del self._systems + self.__systems.extend(self.__final_blocks) # Recurrent call finalize functions for all components. for finalize in self._feature_group_finalize: diff --git a/elastica/modules/callbacks.py b/elastica/modules/callbacks.py index 4a299da0a..de1e50918 100644 --- a/elastica/modules/callbacks.py +++ b/elastica/modules/callbacks.py @@ -51,7 +51,7 @@ def collect_diagnostics( ------- """ - sys_idx: SystemIdxType = self._get_sys_idx_if_valid(system) + sys_idx: SystemIdxType = self.get_system_index(system) # Create _Constraint object, cache it and return to user _callbacks: ModuleProtocol = _CallBack(sys_idx) @@ -77,7 +77,7 @@ def _callback_execution( current_step: int, ) -> None: for sys_id, callback in self._callback_operators: - callback.make_callback(self._systems[sys_id], time, current_step) + callback.make_callback(self[sys_id], time, current_step) class _CallBack: diff --git a/elastica/modules/connections.py b/elastica/modules/connections.py index c55d5beaf..aa1f63c8c 100644 --- a/elastica/modules/connections.py +++ b/elastica/modules/connections.py @@ -66,8 +66,8 @@ def connect( """ # For each system identified, get max dofs - sys_idx_first = self._get_sys_idx_if_valid(first_rod) - sys_idx_second = self._get_sys_idx_if_valid(second_rod) + sys_idx_first = self.get_system_index(first_rod) + sys_idx_second = self.get_system_index(second_rod) sys_dofs_first = first_rod.n_elems sys_dofs_second = second_rod.n_elems @@ -118,9 +118,9 @@ def apply_forces_and_torques( func = functools.partial( apply_forces_and_torques, connect_instance=connect_instance, - system_one=self._systems[first_sys_idx], + system_one=self[first_sys_idx], first_connect_idx=first_connect_idx, - system_two=self._systems[second_sys_idx], + system_two=self[second_sys_idx], second_connect_idx=second_connect_idx, ) diff --git a/elastica/modules/constraints.py b/elastica/modules/constraints.py index 4a2c58407..029ed9615 100644 --- a/elastica/modules/constraints.py +++ b/elastica/modules/constraints.py @@ -57,7 +57,7 @@ def constrain( ------- """ - sys_idx = self._get_sys_idx_if_valid(system) + sys_idx = self.get_system_index(system) # Create _Constraint object, cache it and return to user _constraint: ModuleProtocol = _Constraint(sys_idx) @@ -73,13 +73,13 @@ def _finalize_constraints(self: SystemCollectionProtocol) -> None: """ from elastica._synchronize_periodic_boundary import _ConstrainPeriodicBoundaries - for block in self.systems(): + for block in self.block_systems(): # append the memory block to the simulation as a system. Memory block is the final system in the simulation. if hasattr(block, "ring_rod_flag"): # Apply the constrain to synchronize the periodic boundaries of the memory rod. Find the memory block # sys idx among other systems added and then apply boundary conditions. - memory_block_idx = self._get_sys_idx_if_valid(block) - block_system = cast(BlockSystemType, self._systems[memory_block_idx]) + memory_block_idx = self.get_system_index(block) + block_system = cast(BlockSystemType, self[memory_block_idx]) self.constrain(block_system).using( _ConstrainPeriodicBoundaries, ) @@ -90,7 +90,7 @@ def _finalize_constraints(self: SystemCollectionProtocol) -> None: # dev : the first index stores the rod index to apply the boundary condition # to. self._constraints_operators = [ - (constraint.id(), constraint.instantiate(self._systems[constraint.id()])) + (constraint.id(), constraint.instantiate(self[constraint.id()])) for constraint in self._constraints_list ] @@ -109,11 +109,11 @@ def _finalize_constraints(self: SystemCollectionProtocol) -> None: def _constrain_values(self: SystemCollectionProtocol, time: np.float64) -> None: for sys_id, constraint in self._constraints_operators: - constraint.constrain_values(self._systems[sys_id], time) + constraint.constrain_values(self[sys_id], time) def _constrain_rates(self: SystemCollectionProtocol, time: np.float64) -> None: for sys_id, constraint in self._constraints_operators: - constraint.constrain_rates(self._systems[sys_id], time) + constraint.constrain_rates(self[sys_id], time) class _Constraint: diff --git a/elastica/modules/contact.py b/elastica/modules/contact.py index d30136355..764916644 100644 --- a/elastica/modules/contact.py +++ b/elastica/modules/contact.py @@ -63,8 +63,8 @@ def detect_contact_between( ------- """ - sys_idx_first = self._get_sys_idx_if_valid(first_system) - sys_idx_second = self._get_sys_idx_if_valid(second_system) + sys_idx_first = self.get_system_index(first_system) + sys_idx_second = self.get_system_index(second_system) # Create _Contact object, cache it and return to user _contact = _Contact(sys_idx_first, sys_idx_second) @@ -86,8 +86,8 @@ def apply_contact( second_sys_idx: SystemIdxType, ) -> None: contact_instance.apply_contact( - system_one=self._systems[first_sys_idx], - system_two=self._systems[second_sys_idx], + system_one=self[first_sys_idx], + system_two=self[second_sys_idx], ) for contact in self._contacts: @@ -95,8 +95,8 @@ def apply_contact( contact_instance = contact.instantiate() contact_instance._check_systems_validity( - self._systems[first_sys_idx], - self._systems[second_sys_idx], + self[first_sys_idx], + self[second_sys_idx], ) func = functools.partial( apply_contact, diff --git a/elastica/modules/damping.py b/elastica/modules/damping.py index 646e1a21d..ea9ea5ea6 100644 --- a/elastica/modules/damping.py +++ b/elastica/modules/damping.py @@ -52,7 +52,7 @@ def dampen(self: SystemCollectionProtocol, system: RodType) -> ModuleProtocol: ------- """ - sys_idx = self._get_sys_idx_if_valid(system) + sys_idx = self.get_system_index(system) # Create _Damper object, cache it and return to user _damper: ModuleProtocol = _Damper(sys_idx) @@ -65,7 +65,7 @@ def _finalize_dampers(self: SystemCollectionProtocol) -> None: # inplace : https://stackoverflow.com/a/1208792 self._damping_operators = [ - (damper.id(), damper.instantiate(self._systems[damper.id()])) + (damper.id(), damper.instantiate(self[damper.id()])) for damper in self._damping_list ] @@ -78,7 +78,7 @@ def _finalize_dampers(self: SystemCollectionProtocol) -> None: def _dampen_rates(self: SystemCollectionProtocol, time: np.float64) -> None: for sys_id, damper in self._damping_operators: - damper.dampen_rates(self._systems[sys_id], time) + damper.dampen_rates(self[sys_id], time) class _Damper: diff --git a/elastica/modules/forcing.py b/elastica/modules/forcing.py index d6475d8a2..e36b08999 100644 --- a/elastica/modules/forcing.py +++ b/elastica/modules/forcing.py @@ -53,7 +53,7 @@ def add_forcing_to( ------- """ - sys_idx = self._get_sys_idx_if_valid(system) + sys_idx = self.get_system_index(system) # Create _Constraint object, cache it and return to user _ext_force_torque = _ExtForceTorque(sys_idx) @@ -73,10 +73,10 @@ def _finalize_forcing(self: SystemCollectionProtocol) -> None: forcing_instance = external_force_and_torque.instantiate() apply_forces = functools.partial( - forcing_instance.apply_forces, system=self._systems[sys_id] + forcing_instance.apply_forces, system=self[sys_id] ) apply_torques = functools.partial( - forcing_instance.apply_torques, system=self._systems[sys_id] + forcing_instance.apply_torques, system=self[sys_id] ) self._feature_group_synchronize.add_operators( diff --git a/elastica/modules/memory_block.py b/elastica/modules/memory_block.py index d0e751c15..fa75f744b 100644 --- a/elastica/modules/memory_block.py +++ b/elastica/modules/memory_block.py @@ -8,7 +8,6 @@ RigidBodyType, SurfaceType, StaticSystemType, - SystemType, SystemIdxType, BlockSystemType, ) @@ -22,7 +21,7 @@ def construct_memory_block_structures( systems: list[StaticSystemType], -) -> list[SystemType]: +) -> list[BlockSystemType]: """ This function takes the systems (rod or rigid body) appended to the simulator class and separates them into lists depending on if system is Cosserat rod or rigid body. Then using diff --git a/elastica/modules/protocol.py b/elastica/modules/protocol.py index 4ca93a17c..eb1b5d842 100644 --- a/elastica/modules/protocol.py +++ b/elastica/modules/protocol.py @@ -10,6 +10,7 @@ OperatorFinalizeType, StaticSystemType, SystemType, + BlockSystemType, ConnectionIndex, ) from elastica.joint import FreeJoint @@ -34,10 +35,12 @@ def id(self) -> Any: ... class SystemCollectionProtocol(Protocol): - _systems: list[StaticSystemType] - def __len__(self) -> int: ... + def systems(self) -> Generator[StaticSystemType, None, None]: ... + + def block_systems(self) -> Generator[BlockSystemType, None, None]: ... + @overload def __getitem__(self, i: slice) -> list[SystemType]: ... @overload @@ -67,9 +70,7 @@ def apply_callbacks(self, time: np.float64, current_step: int) -> None: ... @property def _feature_group_finalize(self) -> list[OperatorFinalizeType]: ... - def systems(self) -> Generator[SystemType, None, None]: ... - - def _get_sys_idx_if_valid( + def get_system_index( self, sys_to_be_added: "SystemType | StaticSystemType" ) -> SystemIdxType: ... diff --git a/elastica/systems/analytical.py b/elastica/systems/analytical.py index 317b3f99e..60fe712b9 100644 --- a/elastica/systems/analytical.py +++ b/elastica/systems/analytical.py @@ -304,6 +304,9 @@ def __init__(self): def systems(self): return self._memory_blocks + def block_systems(self): + return self._memory_blocks + def __getitem__(self, idx): return self._memory_blocks[idx] diff --git a/elastica/timestepper/symplectic_steppers.py b/elastica/timestepper/symplectic_steppers.py index 29c3a7e06..3937e337b 100644 --- a/elastica/timestepper/symplectic_steppers.py +++ b/elastica/timestepper/symplectic_steppers.py @@ -97,7 +97,7 @@ def do_step( """ for kin_prefactor, kin_step, dyn_step in steps_and_prefactors[:-1]: - for system in SystemCollection.systems(): + for system in SystemCollection.block_systems(): kin_step(system, time, dt) time += kin_prefactor(dt) @@ -106,14 +106,14 @@ def do_step( SystemCollection.constrain_values(time) # We need internal forces and torques because they are used by interaction module. - for system in SystemCollection.systems(): + for system in SystemCollection.block_systems(): system.compute_internal_forces_and_torques(time) # system.update_internal_forces_and_torques() # Add external forces, controls etc. SystemCollection.synchronize(time) - for system in SystemCollection.systems(): + for system in SystemCollection.block_systems(): dyn_step(system, time, dt) # Constrain only rates @@ -123,7 +123,7 @@ def do_step( last_kin_prefactor = steps_and_prefactors[-1][0] last_kin_step = steps_and_prefactors[-1][1] - for system in SystemCollection.systems(): + for system in SystemCollection.block_systems(): last_kin_step(system, time, dt) time += last_kin_prefactor(dt) SystemCollection.constrain_values(time) @@ -132,7 +132,7 @@ def do_step( SystemCollection.apply_callbacks(time, round(time / dt)) # Zero out the external forces and torques - for system in SystemCollection.systems(): + for system in SystemCollection.block_systems(): system.zeroed_out_external_forces_and_torques(time) return time diff --git a/tests/test_modules/test_base_system.py b/tests/test_modules/test_base_system.py index e3a74dcd8..c126fec4e 100644 --- a/tests/test_modules/test_base_system.py +++ b/tests/test_modules/test_base_system.py @@ -25,11 +25,12 @@ def test_check_type_with_illegal_type_throws(self, illegal_type): @pytest.fixture(scope="class") def load_collection(self): bsc = BaseSystemCollection() + bsc.extend_allowed_types((int, float, str, np.ndarray)) # Bypass check, but its fine for testing - bsc._systems.append(3) - bsc._systems.append(5.0) - bsc._systems.append("a") - bsc._systems.append(np.random.randn(3, 5)) + bsc.append(3) + bsc.append(5.0) + bsc.append("a") + bsc.append(np.random.randn(3, 5)) return bsc def test_len(self, load_collection): @@ -70,12 +71,12 @@ def test_str(self, load_collection): def test_extend_allowed_types(self, load_collection): bsc = load_collection - bsc.extend_allowed_types((int, float, str)) from elastica.rod import RodBase from elastica.rigidbody import RigidBodyBase from elastica.surface import SurfaceBase + # Types are extended in the fixture assert bsc.allowed_sys_types == ( RodBase, RigidBodyBase, @@ -83,6 +84,7 @@ def test_extend_allowed_types(self, load_collection): int, float, str, + np.ndarray, ) def test_extend_correctness(self, load_collection): @@ -121,11 +123,11 @@ def test_invalid_idx_in_get_sys_index_throws(self, load_collection): bsc = load_collection bsc.override_allowed_types((RodBase,)) with pytest.raises(AssertionError) as excinfo: - bsc._get_sys_idx_if_valid(100) + bsc.get_system_index(100) assert "exceeds number of" in str(excinfo.value) with pytest.raises(AssertionError) as excinfo: - load_collection._get_sys_idx_if_valid(np.int_(100)) + load_collection.get_system_index(np.int_(100)) assert "exceeds number of" in str(excinfo.value) def test_unregistered_system_in_get_sys_index_throws( @@ -135,11 +137,11 @@ def test_unregistered_system_in_get_sys_index_throws( my_mock_rod = mock_rod with pytest.raises(ValueError) as excinfo: - load_collection._get_sys_idx_if_valid(my_mock_rod) + load_collection.get_system_index(my_mock_rod) assert "was not found, did you" in str(excinfo.value) def test_get_sys_index_returns_correct_idx(self, load_collection): - assert load_collection._get_sys_idx_if_valid(1) == 1 + assert load_collection.get_system_index(1) == 1 @pytest.mark.xfail def test_delitem(self, load_collection): @@ -171,7 +173,7 @@ def load_collection(self): youngs_modulus=1, ) # Bypass check, but its fine for testing - sc._systems.append(rod) + sc.append(rod) return sc, rod diff --git a/tests/test_modules/test_callbacks.py b/tests/test_modules/test_callbacks.py index b6216cc5f..1716d2fad 100644 --- a/tests/test_modules/test_callbacks.py +++ b/tests/test_modules/test_callbacks.py @@ -80,7 +80,7 @@ def load_system_with_callbacks(self, request): sys_coll_with_callbacks.append(self.MockRod(2, 3, 4, 5)) return sys_coll_with_callbacks - """ The following calls test _get_sys_idx_if_valid from BaseSystem indirectly, + """ The following calls test get_system_index from BaseSystem indirectly, and are here because of legacy reasons. I have not removed them because there are Callbacks require testing against multiple indices, which is still use ful to cross-verify against. diff --git a/tests/test_modules/test_connections.py b/tests/test_modules/test_connections.py index 53ed22f46..26f0cb24c 100644 --- a/tests/test_modules/test_connections.py +++ b/tests/test_modules/test_connections.py @@ -201,7 +201,7 @@ def load_system_with_connects(self, request): sys_coll_with_connects.append(self.MockRod(2, 3, 4, 5)) return sys_coll_with_connects - """ The following calls test _get_sys_idx_if_valid from BaseSystem indirectly, + """ The following calls test get_system_index from BaseSystem indirectly, and are here because of legacy reasons. I have not removed them because there are Connections require testing against multiple indices, which is still use ful to cross-verify against. @@ -401,20 +401,20 @@ def test_connect_call_on_systems(self, load_rod_with_connects_and_indices): connect = connection.instantiate() end_distance_vector = ( - system_collection_with_connections_and_indices._systems[ + system_collection_with_connections_and_indices[ sidx ].position_collection[..., sconnect] - - system_collection_with_connections_and_indices._systems[ + - system_collection_with_connections_and_indices[ fidx ].position_collection[..., fconnect] ) elastic_force = connect.k * end_distance_vector relative_velocity = ( - system_collection_with_connections_and_indices._systems[ + system_collection_with_connections_and_indices[ sidx ].velocity_collection[..., sconnect] - - system_collection_with_connections_and_indices._systems[ + - system_collection_with_connections_and_indices[ fidx ].velocity_collection[..., fconnect] ) @@ -423,16 +423,16 @@ def test_connect_call_on_systems(self, load_rod_with_connects_and_indices): contact_force = elastic_force + damping_force assert_allclose( - system_collection_with_connections_and_indices._systems[ - fidx - ].external_forces[..., fconnect], + system_collection_with_connections_and_indices[fidx].external_forces[ + ..., fconnect + ], contact_force, atol=Tolerance.atol(), ) assert_allclose( - system_collection_with_connections_and_indices._systems[ - sidx - ].external_forces[..., sconnect], + system_collection_with_connections_and_indices[sidx].external_forces[ + ..., sconnect + ], -1 * contact_force, atol=Tolerance.atol(), ) diff --git a/tests/test_modules/test_constraints.py b/tests/test_modules/test_constraints.py index 069479d1a..e5e5892ec 100644 --- a/tests/test_modules/test_constraints.py +++ b/tests/test_modules/test_constraints.py @@ -227,7 +227,7 @@ def load_system_with_constraints(self, request): sys_coll_with_constraints.append(self.MockRod(2, 3, 4, 5)) return sys_coll_with_constraints - """ The following calls test _get_sys_idx_if_valid from BaseSystem indirectly, + """ The following calls test get_system_index from BaseSystem indirectly, and are here because of legacy reasons. I have not removed them because there are Connections require testing against multiple indices, which is still use ful to cross-verify against. @@ -328,11 +328,11 @@ def test_constraint_properties(self, load_rod_with_constraints): for i in [0, 1, -1]: x, y = scwc._constraints_operators[i] - mock_rod = scwc._systems[i] + mock_rod = scwc[i] # Test system assert type(x) is int assert type(y.system) is type(mock_rod) - assert y.system is mock_rod, f"{len(scwc._systems)}" + assert y.system is mock_rod, f"{len(scwc)}" # Test node indices assert y.constrained_position_idx.size == 0 # Test element indices. TODO: maybe add more generalized test diff --git a/tests/test_modules/test_contact.py b/tests/test_modules/test_contact.py index 5b9ffe75b..19ea35608 100644 --- a/tests/test_modules/test_contact.py +++ b/tests/test_modules/test_contact.py @@ -125,7 +125,7 @@ def load_system_with_contacts(self, request): sys_coll_with_contacts.append(self.MockRod(2, 3, 4, 5)) return sys_coll_with_contacts - """ The following calls test _get_sys_idx_if_valid from BaseSystem indirectly, + """ The following calls test get_system_index from BaseSystem indirectly, and are here because of legacy reasons. I have not removed them because there are Contacts require testing against multiple indices, which is still use ful to cross-verify against. @@ -364,8 +364,8 @@ def test_contact_call_on_systems(self, load_system_with_rods_in_contact): fidx, sidx = _contact.id() contact = _contact.instantiate() - system_one = system_collection_with_rods_in_contact._systems[fidx] - system_two = system_collection_with_rods_in_contact._systems[sidx] + system_one = system_collection_with_rods_in_contact[fidx] + system_two = system_collection_with_rods_in_contact[sidx] external_forces_system_one = np.zeros_like(system_one.external_forces) external_forces_system_two = np.zeros_like(system_two.external_forces) @@ -393,12 +393,12 @@ def test_contact_call_on_systems(self, load_system_with_rods_in_contact): ) assert_allclose( - system_collection_with_rods_in_contact._systems[fidx].external_forces, + system_collection_with_rods_in_contact[fidx].external_forces, external_forces_system_one, atol=Tolerance.atol(), ) assert_allclose( - system_collection_with_rods_in_contact._systems[sidx].external_forces, + system_collection_with_rods_in_contact[sidx].external_forces, external_forces_system_two, atol=Tolerance.atol(), ) diff --git a/tests/test_modules/test_damping.py b/tests/test_modules/test_damping.py index 2f18a5d64..70ae6cf70 100644 --- a/tests/test_modules/test_damping.py +++ b/tests/test_modules/test_damping.py @@ -103,7 +103,7 @@ def load_system_with_dampers(self, request): sys_coll_with_dampers.append(self.MockRod(2, 3, 4, 5)) return sys_coll_with_dampers - """ The following calls test _get_sys_idx_if_valid from BaseSystem indirectly, + """ The following calls test get_system_index from BaseSystem indirectly, and are here because of legacy reasons. I have not removed them because there are Connections require testing against multiple indices, which is still use ful to cross-verify against. @@ -195,11 +195,11 @@ def test_damper_properties(self, load_rod_with_dampers): for i in [0, 1, -1]: x, y = scwd._damping_operators[i] - mock_rod = scwd._systems[i] + mock_rod = scwd[i] # Test system assert type(x) is int assert type(y.system) is type(mock_rod) - assert y.system is mock_rod, f"{len(scwd._systems)}" + assert y.system is mock_rod, f"{len(scwd)}" @pytest.mark.xfail def test_dampers_finalize_sorted(self, load_rod_with_dampers): diff --git a/tests/test_modules/test_forcing.py b/tests/test_modules/test_forcing.py index bd384fc6c..74cb2bf3b 100644 --- a/tests/test_modules/test_forcing.py +++ b/tests/test_modules/test_forcing.py @@ -87,7 +87,7 @@ def load_system_with_forcings(self, request): sys_coll_with_forcings.append(self.MockRod(2, 3, 4, 5)) return sys_coll_with_forcings - """ The following calls test _get_sys_idx_if_valid from BaseSystem indirectly, + """ The following calls test get_system_index from BaseSystem indirectly, and are here because of legacy reasons. I have not removed them because there are Connections require testing against multiple indices, which is still use ful to cross-verify against. diff --git a/tests/test_restart.py b/tests/test_restart.py index f19cf0377..28d305957 100644 --- a/tests/test_restart.py +++ b/tests/test_restart.py @@ -41,7 +41,7 @@ def load_collection(self): youngs_modulus=1, ) # Bypass check, but its fine for testing - sc._systems.append(rod) + sc.append(rod) # Also add rods to a separate list rod_list.append(rod) @@ -100,7 +100,7 @@ class BaseSimulatorClass( youngs_modulus=1, ) # Bypass check, but its fine for testing - simulator_class._systems.append(rod) + simulator_class.append(rod) # Also add rods to a separate list rod_list.append(rod) @@ -199,7 +199,7 @@ def load_collection(self): density=1, ) # Bypass check, but its fine for testing - sc._systems.append(cylinder) + sc.append(cylinder) # Also add rods to a separate list cylinder_list.append(cylinder) From 0bd4705cd2e3373911657f7449550b7b6621fe17 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Wed, 26 Jun 2024 00:05:16 -0500 Subject: [PATCH 131/134] type: remove None-type for connection index for simplicity --- elastica/modules/base_system.py | 47 ++++++++++++++++++++++---- elastica/modules/connections.py | 23 +++++++------ elastica/typing.py | 5 +-- tests/test_modules/test_connections.py | 5 +-- 4 files changed, 59 insertions(+), 21 deletions(-) diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index 43666bf66..a7978ca9d 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -144,6 +144,24 @@ def override_allowed_types( def get_system_index( self, system: "SystemType | StaticSystemType" ) -> SystemIdxType: + """ + Get the index of the system object in the system list. + System list is private, so this is the only way to get the index of the system object. + + Example + ------- + >>> system_collection: SystemCollectionProtocol + >>> system: SystemType + ... + >>> system_idx = system_collection.get_system_index(system) # save idx + ... + >>> system = system_collection[system_idx] # just need idx to retrieve + + Parameters + ---------- + system: SystemType + System object to be found in the system list. + """ n_systems = len(self) # Total number of systems from mixed-in class sys_idx: SystemIdxType @@ -170,13 +188,18 @@ def get_system_index( @final def systems(self) -> Generator[StaticSystemType, None, None]: - # assert self._finalize_flag, "The simulator is not finalized." + """ + Iterate over all systems in the system collection. + If the system collection is finalized, block objects are also included. + """ for system in self.__systems: yield system @final def block_systems(self) -> Generator[BlockSystemType, None, None]: - # assert self._finalize_flag, "The simulator is not finalized." + """ + Iterate over all block systems in the system collection. + """ for block in self.__final_blocks: yield block @@ -208,24 +231,36 @@ def finalize(self) -> None: @final def synchronize(self, time: np.float64) -> None: - # Collection call _feature_group_synchronize + """ + Call synchronize functions for all features. + Features are registered in _feature_group_synchronize. + """ for func in self._feature_group_synchronize: func(time=time) @final def constrain_values(self, time: np.float64) -> None: - # Collection call _feature_group_constrain_values + """ + Call constrain values functions for all features. + Features are registered in _feature_group_constrain_values. + """ for func in self._feature_group_constrain_values: func(time=time) @final def constrain_rates(self, time: np.float64) -> None: - # Collection call _feature_group_constrain_rates + """ + Call constrain rates functions for all features. + Features are registered in _feature_group_constrain_rates. + """ for func in self._feature_group_constrain_rates: func(time=time) @final def apply_callbacks(self, time: np.float64, current_step: int) -> None: - # Collection call _feature_group_callback + """ + Call callback functions for all features. + Features are registered in _feature_group_callback. + """ for func in self._feature_group_callback: func(time=time, current_step=current_step) diff --git a/elastica/modules/connections.py b/elastica/modules/connections.py index aa1f63c8c..99f1c9ae7 100644 --- a/elastica/modules/connections.py +++ b/elastica/modules/connections.py @@ -42,8 +42,8 @@ def connect( self: SystemCollectionProtocol, first_rod: "RodType | RigidBodyType", second_rod: "RodType | RigidBodyType", - first_connect_idx: ConnectionIndex = None, - second_connect_idx: ConnectionIndex = None, + first_connect_idx: ConnectionIndex = (), + second_connect_idx: ConnectionIndex = (), ) -> ModuleProtocol: """ This method connects two rod-like objects using the selected joint class. @@ -56,9 +56,9 @@ def connect( Rod-like object second_rod : RodType | RigidBodyType Rod-like object - first_connect_idx : Optional[int] + first_connect_idx : ConnectionIndex Index of first rod for joint. - second_connect_idx : Optional[int] + second_connect_idx : ConnectionIndex Index of second rod for joint. Returns @@ -173,8 +173,8 @@ def __init__( self._second_sys_idx: SystemIdxType = second_sys_idx self._first_sys_n_lim: int = first_sys_nlim self._second_sys_n_lim: int = second_sys_nlim - self.first_sys_connection_idx: ConnectionIndex = None - self.second_sys_connection_idx: ConnectionIndex = None + self.first_sys_connection_idx: ConnectionIndex = () + self.second_sys_connection_idx: ConnectionIndex = () self._connect_cls: Type[FreeJoint] def set_index( @@ -188,7 +188,7 @@ def set_index( ), f"Type of first_connect_idx :{first_type} is different than second_connect_idx :{second_type}" # Check if the type of idx variables are correct. - allow_types = (int, np.int_, list, tuple, np.ndarray, type(None)) + allow_types = (int, np.int_, list, tuple, np.ndarray) assert isinstance( first_idx, allow_types ), f"Connection index type is not supported :{first_type}, please try one of the following :{allow_types}" @@ -227,10 +227,7 @@ def set_index( ), "Connection index of second rod exceeds its dof : {}".format( self._second_sys_n_lim ) - elif first_idx is None: - # Do nothing if idx are None - pass - else: + elif isinstance(first_idx, (int, np.int_)): # The addition of +1 and and <= check on the RHS is because # connections can be made to the node indices as well first_idx__ = cast(int, first_idx) @@ -245,6 +242,10 @@ def set_index( ), "Connection index of second rod exceeds its dof : {}".format( self._second_sys_n_lim ) + else: + raise TypeError( + "Connection index type is not supported :{}".format(first_type) + ) self.first_sys_connection_idx = first_idx self.second_sys_connection_idx = second_idx diff --git a/elastica/typing.py b/elastica/typing.py index 1ba1a4ee3..d05c4feab 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -57,12 +57,13 @@ # Indexing types # TODO: Maybe just use slice?? -ConstrainingIndex: TypeAlias = list[int] | tuple[int, ...] | np.typing.NDArray +ConstrainingIndex: TypeAlias = tuple[int, ...] ConnectionIndex: TypeAlias = ( - int | np.int_ | list[int] | tuple[int] | np.typing.NDArray | None + int | np.int_ | list[int] | tuple[int, ...] | np.typing.NDArray[np.int32] ) # Operators in elastica.modules +# TODO: can be more specific. OperatorParam = ParamSpec("OperatorParam") OperatorCallbackType: TypeAlias = Callable[..., None] OperatorFinalizeType: TypeAlias = Callable[..., None] diff --git a/tests/test_modules/test_connections.py b/tests/test_modules/test_connections.py index 26f0cb24c..c63f8bfbf 100644 --- a/tests/test_modules/test_connections.py +++ b/tests/test_modules/test_connections.py @@ -108,7 +108,7 @@ def test_set_index_with_illegal_type_second_idx_throws( # Below test is to increase code coverage. If we pass nothing or idx=None, then do nothing. def test_set_index_no_input(self, load_connect): - load_connect.set_index(first_idx=None, second_idx=None) + load_connect.set_index(first_idx=(), second_idx=()) @pytest.mark.parametrize( "legal_idx", [(80, 80), (0, 50), (50, 0), (-20, -20), (-20, 50), (-50, -20)] @@ -291,7 +291,8 @@ def test_connect_registers_and_returns_Connect(self, load_system_with_connects): assert _mock_connect in system_collection_with_connections._connections assert _mock_connect.__class__ == _Connect # check sane defaults provided for connection indices - assert _mock_connect.id()[2] is None and _mock_connect.id()[3] is None + assert _mock_connect.id()[2] == () + assert _mock_connect.id()[3] == () from elastica.joint import FreeJoint From 901b19ac9b736112cec1a3daccb5a364326a5489 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Wed, 26 Jun 2024 08:52:42 -0500 Subject: [PATCH 132/134] remove out-dated function --- elastica/_calculus.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/elastica/_calculus.py b/elastica/_calculus.py index c6704480b..3cdda50d9 100644 --- a/elastica/_calculus.py +++ b/elastica/_calculus.py @@ -1,5 +1,4 @@ __doc__ = """ Quadrature and difference kernels """ -from typing import Union import numpy as np from numpy import zeros, empty from numpy.typing import NDArray @@ -8,17 +7,6 @@ from elastica.reset_functions_for_block_structure._reset_ghost_vector_or_scalar import ( _reset_vector_ghost, ) -import functools - - -@functools.lru_cache(maxsize=2) -def _get_zero_array(dim: int, ndim: int) -> Union[float, NDArray[np.float64], None]: - if ndim == 1: - return 0.0 - if ndim == 2: - return np.zeros((dim, 1)) - - return None @njit(cache=True) # type: ignore From e04358a52c3f6a88c832f1f073d00b7ee64d7d31 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Wed, 26 Jun 2024 09:21:53 -0500 Subject: [PATCH 133/134] fix typo --- elastica/modules/base_system.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index a7978ca9d..9f8415dc3 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -37,9 +37,9 @@ class BaseSystemCollection(MutableSequence): Attributes ---------- - allowed_sys_types: tuple + allowed_sys_types: tuple[Type] Tuple of allowed type rod-like objects. Here use a base class for objects, i.e. RodBase. - systems: Callabke + systems: Callable Returns all system objects. Once finalize, block objects are also included. blocks: Callable Returns block objects. Should be called after finalize. From a8d9a6f46d7018a1740afd2baf4a33f31c5237a6 Mon Sep 17 00:00:00 2001 From: Seung Hyun Kim Date: Thu, 27 Jun 2024 10:43:51 -0500 Subject: [PATCH 134/134] change: int_ -> int32 --- elastica/boundary_conditions.py | 8 ++++---- elastica/modules/base_system.py | 4 +++- elastica/modules/connections.py | 14 ++++++++++---- elastica/typing.py | 2 +- tests/test_modules/test_base_system.py | 2 +- tests/test_modules/test_callbacks.py | 2 +- tests/test_modules/test_connections.py | 2 +- tests/test_modules/test_constraints.py | 2 +- tests/test_modules/test_contact.py | 2 +- tests/test_modules/test_damping.py | 2 +- tests/test_modules/test_forcing.py | 2 +- 11 files changed, 25 insertions(+), 17 deletions(-) diff --git a/elastica/boundary_conditions.py b/elastica/boundary_conditions.py index e5ac86b4d..1075821cf 100644 --- a/elastica/boundary_conditions.py +++ b/elastica/boundary_conditions.py @@ -48,10 +48,10 @@ def __init__( try: self._system = kwargs["_system"] self._constrained_position_idx = np.array( - constrained_position_idx, dtype=np.int_ + constrained_position_idx, dtype=np.int32 ) self._constrained_director_idx = np.array( - constrained_director_idx, dtype=np.int_ + constrained_director_idx, dtype=np.int32 ) except KeyError: raise KeyError( @@ -64,12 +64,12 @@ def system(self) -> S: return self._system @property - def constrained_position_idx(self) -> np.ndarray: + def constrained_position_idx(self) -> NDArray[np.int32]: """get position-indices passed to "using" """ return self._constrained_position_idx @property - def constrained_director_idx(self) -> np.ndarray: + def constrained_director_idx(self) -> NDArray[np.int32]: """get director-indices passed to "using" """ return self._constrained_director_idx diff --git a/elastica/modules/base_system.py b/elastica/modules/base_system.py index a7978ca9d..6a8e630cd 100644 --- a/elastica/modules/base_system.py +++ b/elastica/modules/base_system.py @@ -165,7 +165,9 @@ def get_system_index( n_systems = len(self) # Total number of systems from mixed-in class sys_idx: SystemIdxType - if isinstance(system, (int, np.int_)): + if isinstance( + system, (int, np.integer) + ): # np.integer includes both int32 and int64 # 1. If they are indices themselves, check range # This is only used for testing purposes assert ( diff --git a/elastica/modules/connections.py b/elastica/modules/connections.py index 99f1c9ae7..44a590732 100644 --- a/elastica/modules/connections.py +++ b/elastica/modules/connections.py @@ -188,7 +188,13 @@ def set_index( ), f"Type of first_connect_idx :{first_type} is different than second_connect_idx :{second_type}" # Check if the type of idx variables are correct. - allow_types = (int, np.int_, list, tuple, np.ndarray) + allow_types = ( + int, + np.integer, + list, + tuple, + np.ndarray, + ) # np.integer is for both int32 and int64 assert isinstance( first_idx, allow_types ), f"Connection index type is not supported :{first_type}, please try one of the following :{allow_types}" @@ -198,13 +204,13 @@ def set_index( first_idx_ = cast(list[int], first_idx) second_idx_ = cast(list[int], second_idx) for i in range(len(first_idx_)): - assert isinstance(first_idx_[i], (int, np.int_)), ( + assert isinstance(first_idx_[i], (int, np.integer)), ( "Connection index of first rod is not integer :{}".format( first_idx_[i] ) + " It should be : integer. Check your input!" ) - assert isinstance(second_idx_[i], (int, np.int_)), ( + assert isinstance(second_idx_[i], (int, np.integer)), ( "Connection index of second rod is not integer :{}".format( second_idx_[i] ) @@ -227,7 +233,7 @@ def set_index( ), "Connection index of second rod exceeds its dof : {}".format( self._second_sys_n_lim ) - elif isinstance(first_idx, (int, np.int_)): + elif isinstance(first_idx, (int, np.integer)): # The addition of +1 and and <= check on the RHS is because # connections can be made to the node indices as well first_idx__ = cast(int, first_idx) diff --git a/elastica/typing.py b/elastica/typing.py index d05c4feab..79cba0036 100644 --- a/elastica/typing.py +++ b/elastica/typing.py @@ -59,7 +59,7 @@ # TODO: Maybe just use slice?? ConstrainingIndex: TypeAlias = tuple[int, ...] ConnectionIndex: TypeAlias = ( - int | np.int_ | list[int] | tuple[int, ...] | np.typing.NDArray[np.int32] + int | np.int32 | list[int] | tuple[int, ...] | np.typing.NDArray[np.int32] ) # Operators in elastica.modules diff --git a/tests/test_modules/test_base_system.py b/tests/test_modules/test_base_system.py index c126fec4e..694bd84f7 100644 --- a/tests/test_modules/test_base_system.py +++ b/tests/test_modules/test_base_system.py @@ -127,7 +127,7 @@ def test_invalid_idx_in_get_sys_index_throws(self, load_collection): assert "exceeds number of" in str(excinfo.value) with pytest.raises(AssertionError) as excinfo: - load_collection.get_system_index(np.int_(100)) + load_collection.get_system_index(np.int32(100)) assert "exceeds number of" in str(excinfo.value) def test_unregistered_system_in_get_sys_index_throws( diff --git a/tests/test_modules/test_callbacks.py b/tests/test_modules/test_callbacks.py index 1716d2fad..a58901e25 100644 --- a/tests/test_modules/test_callbacks.py +++ b/tests/test_modules/test_callbacks.py @@ -96,7 +96,7 @@ def test_callback_with_illegal_index_throws(self, load_system_with_callbacks): assert "exceeds number of" in str(excinfo.value) with pytest.raises(AssertionError) as excinfo: - scwc.collect_diagnostics(np.int_(100)) + scwc.collect_diagnostics(np.int32(100)) assert "exceeds number of" in str(excinfo.value) def test_callback_with_unregistered_system_throws(self, load_system_with_callbacks): diff --git a/tests/test_modules/test_connections.py b/tests/test_modules/test_connections.py index c63f8bfbf..eadf98a0f 100644 --- a/tests/test_modules/test_connections.py +++ b/tests/test_modules/test_connections.py @@ -232,7 +232,7 @@ def test_connect_with_illegal_index_throws( assert "exceeds number of" in str(excinfo.value) with pytest.raises(AssertionError) as excinfo: - system_collection_with_connections.connect(*[np.int_(x) for x in sys_idx]) + system_collection_with_connections.connect(*[np.int32(x) for x in sys_idx]) assert "exceeds number of" in str(excinfo.value) def test_connect_with_unregistered_system_throws(self, load_system_with_connects): diff --git a/tests/test_modules/test_constraints.py b/tests/test_modules/test_constraints.py index e5e5892ec..c6072c392 100644 --- a/tests/test_modules/test_constraints.py +++ b/tests/test_modules/test_constraints.py @@ -243,7 +243,7 @@ def test_constrain_with_illegal_index_throws(self, load_system_with_constraints) assert "exceeds number of" in str(excinfo.value) with pytest.raises(AssertionError) as excinfo: - scwc.constrain(np.int_(100)) + scwc.constrain(np.int32(100)) assert "exceeds number of" in str(excinfo.value) def test_constrain_with_unregistered_system_throws( diff --git a/tests/test_modules/test_contact.py b/tests/test_modules/test_contact.py index 19ea35608..79abb7c5e 100644 --- a/tests/test_modules/test_contact.py +++ b/tests/test_modules/test_contact.py @@ -157,7 +157,7 @@ def test_contact_with_illegal_index_throws( with pytest.raises(AssertionError) as excinfo: system_collection_with_contacts.detect_contact_between( - *[np.int_(x) for x in sys_idx] + *[np.int32(x) for x in sys_idx] ) assert "exceeds number of" in str(excinfo.value) diff --git a/tests/test_modules/test_damping.py b/tests/test_modules/test_damping.py index 70ae6cf70..2634da841 100644 --- a/tests/test_modules/test_damping.py +++ b/tests/test_modules/test_damping.py @@ -119,7 +119,7 @@ def test_dampen_with_illegal_index_throws(self, load_system_with_dampers): assert "exceeds number of" in str(excinfo.value) with pytest.raises(AssertionError) as excinfo: - scwd.dampen(np.int_(100)) + scwd.dampen(np.int32(100)) assert "exceeds number of" in str(excinfo.value) def test_dampen_with_unregistered_system_throws(self, load_system_with_dampers): diff --git a/tests/test_modules/test_forcing.py b/tests/test_modules/test_forcing.py index 74cb2bf3b..4f621fc40 100644 --- a/tests/test_modules/test_forcing.py +++ b/tests/test_modules/test_forcing.py @@ -103,7 +103,7 @@ def test_constrain_with_illegal_index_throws(self, load_system_with_forcings): assert "exceeds number of" in str(excinfo.value) with pytest.raises(AssertionError) as excinfo: - scwf.add_forcing_to(np.int_(100)) + scwf.add_forcing_to(np.int32(100)) assert "exceeds number of" in str(excinfo.value) def test_constrain_with_unregistered_system_throws(self, load_system_with_forcings):