diff --git a/Makefile b/Makefile
index 5f53a169..217562ba 100644
--- a/Makefile
+++ b/Makefile
@@ -31,11 +31,17 @@ help:
@echo " $(R) tests $(E) \t to run tests with the $(P)pytest$(E) package."
@echo " $(R) type $(E) \t to run type checking with the $(P)mypy$(E) package."
+build:
+ @echo "Re-building wheel and dist"
+ @rm -rf dist
+ @poetry build
+ @echo "Created build is located in the $(C)dist$(E) folder."
+
checklist:
@echo "Here is a small pre-release check-list:"
@echo " - Check you are on a tagged $(P)feature/release$(E) branch (see Gitflow workflow)."
@echo " - Run $(D)poetry version$(E) with the right argument and update the version number in $(C)__init__.py$(E)."
- @echo " - Update the pyhdtoolkit version in the $(C)environment.yml$(E) file."
+ @echo " - Update the pyhdtoolkit version in the $(C)docker/environment.yml$(E) file."
@echo " - Check the $(P)feature/release$(E) branch tag matches this release's package version."
@echo " - After merging and pushing this release from $(P)master$(E) to $(P)origin/master$(E):"
@echo " - Run $(D)poetry build$(E) to create a tarball of the new version."
@@ -51,8 +57,9 @@ clean:
@rm -rf .eggs
@echo "Cleaning up bitecode files and python cache."
@find . -type f -name '*.py[co]' -delete -o -type d -name __pycache__ -delete
- @echo "Cleaning up pytest cache."
+ @echo "Cleaning up pytest cache & test artifacts."
@find . -type d -name '*.pytest_cache' -exec rm -rf {} + -o -type f -name '*.pytest_cache' -exec rm -rf {} +
+ @find . -type f -name 'fc.*' -delete -o -type f -name 'fort.*' -delete
@echo "Cleaning up mypy cache."
@find . -type d -name "*.mypy_cache" -exec rm -rf {} +
@echo "Cleaning up coverage reports."
diff --git a/poetry.lock b/poetry.lock
index dad34b84..fe5ed9bf 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -8,17 +8,16 @@ python-versions = "*"
[[package]]
name = "astroid"
-version = "2.4.2"
+version = "2.5"
description = "An abstract syntax tree for Python with inference support."
category = "dev"
optional = false
-python-versions = ">=3.5"
+python-versions = ">=3.6"
[package.dependencies]
-lazy-object-proxy = ">=1.4.0,<1.5.0"
-six = ">=1.12,<2.0"
+lazy-object-proxy = ">=1.4.0"
typed-ast = {version = ">=1.4.0,<1.5", markers = "implementation_name == \"cpython\" and python_version < \"3.8\""}
-wrapt = ">=1.11,<2.0"
+wrapt = ">=1.11,<1.13"
[[package]]
name = "atomicwrites"
@@ -211,7 +210,7 @@ python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*"
[[package]]
name = "importlib-metadata"
-version = "3.4.0"
+version = "3.6.0"
description = "Read metadata from Python packages"
category = "main"
optional = false
@@ -314,11 +313,11 @@ python-versions = ">=3.6"
[[package]]
name = "lazy-object-proxy"
-version = "1.4.3"
+version = "1.5.2"
description = "A fast and thorough lazy object proxy."
category = "dev"
optional = false
-python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*"
+python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*"
[[package]]
name = "livereload"
@@ -485,7 +484,7 @@ mkdocs-material = ">=5.0.0"
[[package]]
name = "mypy"
-version = "0.800"
+version = "0.812"
description = "Optional static typing for Python"
category = "dev"
optional = false
@@ -675,19 +674,22 @@ python-versions = ">=3.5"
[[package]]
name = "pylint"
-version = "2.6.2"
+version = "2.7.1"
description = "python code static checker"
category = "dev"
optional = false
-python-versions = ">=3.5.*"
+python-versions = "~=3.6"
[package.dependencies]
-astroid = ">=2.4.0,<2.5"
+astroid = "2.5.0"
colorama = {version = "*", markers = "sys_platform == \"win32\""}
isort = ">=4.2.5,<6"
mccabe = ">=0.6,<0.7"
toml = ">=0.7.1"
+[package.extras]
+docs = ["sphinx (>=3.2,<4.0)", "python-docs-theme"]
+
[[package]]
name = "pymdown-extensions"
version = "7.1"
@@ -866,7 +868,7 @@ python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*"
[[package]]
name = "tabulate"
-version = "0.8.8"
+version = "0.8.9"
description = "Pretty-print tabular data"
category = "dev"
optional = false
@@ -1010,8 +1012,8 @@ appdirs = [
{file = "appdirs-1.4.4.tar.gz", hash = "sha256:7d5d0167b2b1ba821647616af46a749d1c653740dd0d2415100fe26e27afdf41"},
]
astroid = [
- {file = "astroid-2.4.2-py3-none-any.whl", hash = "sha256:bc58d83eb610252fd8de6363e39d4f1d0619c894b0ed24603b881c02e64c7386"},
- {file = "astroid-2.4.2.tar.gz", hash = "sha256:2f4078c2a41bf377eea06d71c9d2ba4eb8f6b1af2135bec27bbbb7d8f12bb703"},
+ {file = "astroid-2.5-py3-none-any.whl", hash = "sha256:87ae7f2398b8a0ae5638ddecf9987f081b756e0e9fc071aeebdca525671fc4dc"},
+ {file = "astroid-2.5.tar.gz", hash = "sha256:b31c92f545517dcc452f284bc9c044050862fbe6d93d2b3de4a215a6b384bf0d"},
]
atomicwrites = [
{file = "atomicwrites-1.4.0-py2.py3-none-any.whl", hash = "sha256:6d1784dea7c0c8d4a5172b6c620f40b6e4cbfdf96d783691f2e1302a7b88e197"},
@@ -1174,8 +1176,8 @@ idna = [
{file = "idna-2.10.tar.gz", hash = "sha256:b307872f855b18632ce0c21c5e45be78c0ea7ae4c15c828c20788b26921eb3f6"},
]
importlib-metadata = [
- {file = "importlib_metadata-3.4.0-py3-none-any.whl", hash = "sha256:ace61d5fc652dc280e7b6b4ff732a9c2d40db2c0f92bc6cb74e07b73d53a1771"},
- {file = "importlib_metadata-3.4.0.tar.gz", hash = "sha256:fa5daa4477a7414ae34e95942e4dd07f62adf589143c875c133c1e53c4eff38d"},
+ {file = "importlib_metadata-3.6.0-py3-none-any.whl", hash = "sha256:144acbe6ec34f170c4a362f52c2d03e82db7d012668ba57ff8f1e668905889b0"},
+ {file = "importlib_metadata-3.6.0.tar.gz", hash = "sha256:080d666daa2faedcc2e96c0c87d3cdb76325261c51650016784834733faa7897"},
]
importlib-resources = [
{file = "importlib_resources-5.1.0-py3-none-any.whl", hash = "sha256:885b8eae589179f661c909d699a546cf10d83692553e34dca1bf5eb06f7f6217"},
@@ -1235,27 +1237,30 @@ kiwisolver = [
{file = "kiwisolver-1.3.1.tar.gz", hash = "sha256:950a199911a8d94683a6b10321f9345d5a3a8433ec58b217ace979e18f16e248"},
]
lazy-object-proxy = [
- {file = "lazy-object-proxy-1.4.3.tar.gz", hash = "sha256:f3900e8a5de27447acbf900b4750b0ddfd7ec1ea7fbaf11dfa911141bc522af0"},
- {file = "lazy_object_proxy-1.4.3-cp27-cp27m-macosx_10_13_x86_64.whl", hash = "sha256:a2238e9d1bb71a56cd710611a1614d1194dc10a175c1e08d75e1a7bcc250d442"},
- {file = "lazy_object_proxy-1.4.3-cp27-cp27m-win32.whl", hash = "sha256:efa1909120ce98bbb3777e8b6f92237f5d5c8ea6758efea36a473e1d38f7d3e4"},
- {file = "lazy_object_proxy-1.4.3-cp27-cp27m-win_amd64.whl", hash = "sha256:4677f594e474c91da97f489fea5b7daa17b5517190899cf213697e48d3902f5a"},
- {file = "lazy_object_proxy-1.4.3-cp27-cp27mu-manylinux1_x86_64.whl", hash = "sha256:0c4b206227a8097f05c4dbdd323c50edf81f15db3b8dc064d08c62d37e1a504d"},
- {file = "lazy_object_proxy-1.4.3-cp34-cp34m-manylinux1_x86_64.whl", hash = "sha256:d945239a5639b3ff35b70a88c5f2f491913eb94871780ebfabb2568bd58afc5a"},
- {file = "lazy_object_proxy-1.4.3-cp34-cp34m-win32.whl", hash = "sha256:9651375199045a358eb6741df3e02a651e0330be090b3bc79f6d0de31a80ec3e"},
- {file = "lazy_object_proxy-1.4.3-cp34-cp34m-win_amd64.whl", hash = "sha256:eba7011090323c1dadf18b3b689845fd96a61ba0a1dfbd7f24b921398affc357"},
- {file = "lazy_object_proxy-1.4.3-cp35-cp35m-manylinux1_x86_64.whl", hash = "sha256:48dab84ebd4831077b150572aec802f303117c8cc5c871e182447281ebf3ac50"},
- {file = "lazy_object_proxy-1.4.3-cp35-cp35m-win32.whl", hash = "sha256:ca0a928a3ddbc5725be2dd1cf895ec0a254798915fb3a36af0964a0a4149e3db"},
- {file = "lazy_object_proxy-1.4.3-cp35-cp35m-win_amd64.whl", hash = "sha256:194d092e6f246b906e8f70884e620e459fc54db3259e60cf69a4d66c3fda3449"},
- {file = "lazy_object_proxy-1.4.3-cp36-cp36m-manylinux1_x86_64.whl", hash = "sha256:97bb5884f6f1cdce0099f86b907aa41c970c3c672ac8b9c8352789e103cf3156"},
- {file = "lazy_object_proxy-1.4.3-cp36-cp36m-win32.whl", hash = "sha256:cb2c7c57005a6804ab66f106ceb8482da55f5314b7fcb06551db1edae4ad1531"},
- {file = "lazy_object_proxy-1.4.3-cp36-cp36m-win_amd64.whl", hash = "sha256:8d859b89baf8ef7f8bc6b00aa20316483d67f0b1cbf422f5b4dc56701c8f2ffb"},
- {file = "lazy_object_proxy-1.4.3-cp37-cp37m-macosx_10_13_x86_64.whl", hash = "sha256:1be7e4c9f96948003609aa6c974ae59830a6baecc5376c25c92d7d697e684c08"},
- {file = "lazy_object_proxy-1.4.3-cp37-cp37m-manylinux1_x86_64.whl", hash = "sha256:d74bb8693bf9cf75ac3b47a54d716bbb1a92648d5f781fc799347cfc95952383"},
- {file = "lazy_object_proxy-1.4.3-cp37-cp37m-win32.whl", hash = "sha256:9b15f3f4c0f35727d3a0fba4b770b3c4ebbb1fa907dbcc046a1d2799f3edd142"},
- {file = "lazy_object_proxy-1.4.3-cp37-cp37m-win_amd64.whl", hash = "sha256:9254f4358b9b541e3441b007a0ea0764b9d056afdeafc1a5569eee1cc6c1b9ea"},
- {file = "lazy_object_proxy-1.4.3-cp38-cp38-manylinux1_x86_64.whl", hash = "sha256:a6ae12d08c0bf9909ce12385803a543bfe99b95fe01e752536a60af2b7797c62"},
- {file = "lazy_object_proxy-1.4.3-cp38-cp38-win32.whl", hash = "sha256:5541cada25cd173702dbd99f8e22434105456314462326f06dba3e180f203dfd"},
- {file = "lazy_object_proxy-1.4.3-cp38-cp38-win_amd64.whl", hash = "sha256:59f79fef100b09564bc2df42ea2d8d21a64fdcda64979c0fa3db7bdaabaf6239"},
+ {file = "lazy-object-proxy-1.5.2.tar.gz", hash = "sha256:5944a9b95e97de1980c65f03b79b356f30a43de48682b8bdd90aa5089f0ec1f4"},
+ {file = "lazy_object_proxy-1.5.2-cp27-cp27m-macosx_10_14_x86_64.whl", hash = "sha256:e960e8be509e8d6d618300a6c189555c24efde63e85acaf0b14b2cd1ac743315"},
+ {file = "lazy_object_proxy-1.5.2-cp27-cp27mu-manylinux1_x86_64.whl", hash = "sha256:522b7c94b524389f4a4094c4bf04c2b02228454ddd17c1a9b2801fac1d754871"},
+ {file = "lazy_object_proxy-1.5.2-cp35-cp35m-manylinux1_x86_64.whl", hash = "sha256:3782931963dc89e0e9a0ae4348b44762e868ea280e4f8c233b537852a8996ab9"},
+ {file = "lazy_object_proxy-1.5.2-cp35-cp35m-manylinux2014_aarch64.whl", hash = "sha256:429c4d1862f3fc37cd56304d880f2eae5bd0da83bdef889f3bd66458aac49128"},
+ {file = "lazy_object_proxy-1.5.2-cp35-cp35m-win32.whl", hash = "sha256:cd1bdace1a8762534e9a36c073cd54e97d517a17d69a17985961265be6d22847"},
+ {file = "lazy_object_proxy-1.5.2-cp35-cp35m-win_amd64.whl", hash = "sha256:ddbdcd10eb999d7ab292677f588b658372aadb9a52790f82484a37127a390108"},
+ {file = "lazy_object_proxy-1.5.2-cp36-cp36m-manylinux1_x86_64.whl", hash = "sha256:ecb5dd5990cec6e7f5c9c1124a37cb2c710c6d69b0c1a5c4aa4b35eba0ada068"},
+ {file = "lazy_object_proxy-1.5.2-cp36-cp36m-manylinux2014_aarch64.whl", hash = "sha256:b6577f15d5516d7d209c1a8cde23062c0f10625f19e8dc9fb59268859778d7d7"},
+ {file = "lazy_object_proxy-1.5.2-cp36-cp36m-win32.whl", hash = "sha256:c8fe2d6ff0ff583784039d0255ea7da076efd08507f2be6f68583b0da32e3afb"},
+ {file = "lazy_object_proxy-1.5.2-cp36-cp36m-win_amd64.whl", hash = "sha256:fa5b2dee0e231fa4ad117be114251bdfe6afe39213bd629d43deb117b6a6c40a"},
+ {file = "lazy_object_proxy-1.5.2-cp37-cp37m-manylinux1_x86_64.whl", hash = "sha256:1d33d6f789697f401b75ce08e73b1de567b947740f768376631079290118ad39"},
+ {file = "lazy_object_proxy-1.5.2-cp37-cp37m-manylinux2014_aarch64.whl", hash = "sha256:57fb5c5504ddd45ed420b5b6461a78f58cbb0c1b0cbd9cd5a43ad30a4a3ee4d0"},
+ {file = "lazy_object_proxy-1.5.2-cp37-cp37m-win32.whl", hash = "sha256:e7273c64bccfd9310e9601b8f4511d84730239516bada26a0c9846c9697617ef"},
+ {file = "lazy_object_proxy-1.5.2-cp37-cp37m-win_amd64.whl", hash = "sha256:6f4e5e68b7af950ed7fdb594b3f19a0014a3ace0fedb86acb896e140ffb24302"},
+ {file = "lazy_object_proxy-1.5.2-cp38-cp38-manylinux1_x86_64.whl", hash = "sha256:cadfa2c2cf54d35d13dc8d231253b7985b97d629ab9ca6e7d672c35539d38163"},
+ {file = "lazy_object_proxy-1.5.2-cp38-cp38-manylinux2014_aarch64.whl", hash = "sha256:e7428977763150b4cf83255625a80a23dfdc94d43be7791ce90799d446b4e26f"},
+ {file = "lazy_object_proxy-1.5.2-cp38-cp38-win32.whl", hash = "sha256:2f2de8f8ac0be3e40d17730e0600619d35c78c13a099ea91ef7fb4ad944ce694"},
+ {file = "lazy_object_proxy-1.5.2-cp38-cp38-win_amd64.whl", hash = "sha256:38c3865bd220bd983fcaa9aa11462619e84a71233bafd9c880f7b1cb753ca7fa"},
+ {file = "lazy_object_proxy-1.5.2-cp39-cp39-macosx_10_14_x86_64.whl", hash = "sha256:8a44e9901c0555f95ac401377032f6e6af66d8fc1fbfad77a7a8b1a826e0b93c"},
+ {file = "lazy_object_proxy-1.5.2-cp39-cp39-manylinux1_x86_64.whl", hash = "sha256:fa7fb7973c622b9e725bee1db569d2c2ee64d2f9a089201c5e8185d482c7352d"},
+ {file = "lazy_object_proxy-1.5.2-cp39-cp39-manylinux2014_aarch64.whl", hash = "sha256:71a1ef23f22fa8437974b2d60fedb947c99a957ad625f83f43fd3de70f77f458"},
+ {file = "lazy_object_proxy-1.5.2-cp39-cp39-win32.whl", hash = "sha256:ef3f5e288aa57b73b034ce9c1f1ac753d968f9069cd0742d1d69c698a0167166"},
+ {file = "lazy_object_proxy-1.5.2-cp39-cp39-win_amd64.whl", hash = "sha256:37d9c34b96cca6787fe014aeb651217944a967a5b165e2cacb6b858d2997ab84"},
]
livereload = [
{file = "livereload-2.6.3.tar.gz", hash = "sha256:776f2f865e59fde56490a56bcc6773b6917366bce0c267c60ee8aaf1a0959869"},
@@ -1376,28 +1381,28 @@ mkdocs-material-extensions = [
{file = "mkdocs_material_extensions-1.0.1-py3-none-any.whl", hash = "sha256:d90c807a88348aa6d1805657ec5c0b2d8d609c110e62b9dce4daf7fa981fa338"},
]
mypy = [
- {file = "mypy-0.800-cp35-cp35m-macosx_10_9_x86_64.whl", hash = "sha256:e1c84c65ff6d69fb42958ece5b1255394714e0aac4df5ffe151bc4fe19c7600a"},
- {file = "mypy-0.800-cp35-cp35m-manylinux1_x86_64.whl", hash = "sha256:947126195bfe4709c360e89b40114c6746ae248f04d379dca6f6ab677aa07641"},
- {file = "mypy-0.800-cp35-cp35m-manylinux2010_x86_64.whl", hash = "sha256:b95068a3ce3b50332c40e31a955653be245666a4bc7819d3c8898aa9fb9ea496"},
- {file = "mypy-0.800-cp35-cp35m-win_amd64.whl", hash = "sha256:ca7ad5aed210841f1e77f5f2f7d725b62c78fa77519312042c719ed2ab937876"},
- {file = "mypy-0.800-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:e32b7b282c4ed4e378bba8b8dfa08e1cfa6f6574067ef22f86bee5b1039de0c9"},
- {file = "mypy-0.800-cp36-cp36m-manylinux1_x86_64.whl", hash = "sha256:e497a544391f733eca922fdcb326d19e894789cd4ff61d48b4b195776476c5cf"},
- {file = "mypy-0.800-cp36-cp36m-manylinux2010_x86_64.whl", hash = "sha256:5615785d3e2f4f03ab7697983d82c4b98af5c321614f51b8f1034eb9ebe48363"},
- {file = "mypy-0.800-cp36-cp36m-win_amd64.whl", hash = "sha256:2b216eacca0ec0ee124af9429bfd858d5619a0725ee5f88057e6e076f9eb1a7b"},
- {file = "mypy-0.800-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:e3b8432f8df19e3c11235c4563a7250666dc9aa7cdda58d21b4177b20256ca9f"},
- {file = "mypy-0.800-cp37-cp37m-manylinux1_x86_64.whl", hash = "sha256:d16c54b0dffb861dc6318a8730952265876d90c5101085a4bc56913e8521ba19"},
- {file = "mypy-0.800-cp37-cp37m-manylinux2010_x86_64.whl", hash = "sha256:0d2fc8beb99cd88f2d7e20d69131353053fbecea17904ee6f0348759302c52fa"},
- {file = "mypy-0.800-cp37-cp37m-win_amd64.whl", hash = "sha256:aa9d4901f3ee1a986a3a79fe079ffbf7f999478c281376f48faa31daaa814e86"},
- {file = "mypy-0.800-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:319ee5c248a7c3f94477f92a729b7ab06bf8a6d04447ef3aa8c9ba2aa47c6dcf"},
- {file = "mypy-0.800-cp38-cp38-manylinux1_x86_64.whl", hash = "sha256:74f5aa50d0866bc6fb8e213441c41e466c86678c800700b87b012ed11c0a13e0"},
- {file = "mypy-0.800-cp38-cp38-manylinux2010_x86_64.whl", hash = "sha256:a301da58d566aca05f8f449403c710c50a9860782148332322decf73a603280b"},
- {file = "mypy-0.800-cp38-cp38-win_amd64.whl", hash = "sha256:b9150db14a48a8fa114189bfe49baccdff89da8c6639c2717750c7ae62316738"},
- {file = "mypy-0.800-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:f5fdf935a46aa20aa937f2478480ebf4be9186e98e49cc3843af9a5795a49a25"},
- {file = "mypy-0.800-cp39-cp39-manylinux1_x86_64.whl", hash = "sha256:6f8425fecd2ba6007e526209bb985ce7f49ed0d2ac1cc1a44f243380a06a84fb"},
- {file = "mypy-0.800-cp39-cp39-manylinux2010_x86_64.whl", hash = "sha256:5ff616787122774f510caeb7b980542a7cc2222be3f00837a304ea85cd56e488"},
- {file = "mypy-0.800-cp39-cp39-win_amd64.whl", hash = "sha256:90b6f46dc2181d74f80617deca611925d7e63007cf416397358aa42efb593e07"},
- {file = "mypy-0.800-py3-none-any.whl", hash = "sha256:3e0c159a7853e3521e3f582adb1f3eac66d0b0639d434278e2867af3a8c62653"},
- {file = "mypy-0.800.tar.gz", hash = "sha256:e0202e37756ed09daf4b0ba64ad2c245d357659e014c3f51d8cd0681ba66940a"},
+ {file = "mypy-0.812-cp35-cp35m-macosx_10_9_x86_64.whl", hash = "sha256:a26f8ec704e5a7423c8824d425086705e381b4f1dfdef6e3a1edab7ba174ec49"},
+ {file = "mypy-0.812-cp35-cp35m-manylinux1_x86_64.whl", hash = "sha256:28fb5479c494b1bab244620685e2eb3c3f988d71fd5d64cc753195e8ed53df7c"},
+ {file = "mypy-0.812-cp35-cp35m-manylinux2010_x86_64.whl", hash = "sha256:9743c91088d396c1a5a3c9978354b61b0382b4e3c440ce83cf77994a43e8c521"},
+ {file = "mypy-0.812-cp35-cp35m-win_amd64.whl", hash = "sha256:d7da2e1d5f558c37d6e8c1246f1aec1e7349e4913d8fb3cb289a35de573fe2eb"},
+ {file = "mypy-0.812-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:4eec37370483331d13514c3f55f446fc5248d6373e7029a29ecb7b7494851e7a"},
+ {file = "mypy-0.812-cp36-cp36m-manylinux1_x86_64.whl", hash = "sha256:d65cc1df038ef55a99e617431f0553cd77763869eebdf9042403e16089fe746c"},
+ {file = "mypy-0.812-cp36-cp36m-manylinux2010_x86_64.whl", hash = "sha256:61a3d5b97955422964be6b3baf05ff2ce7f26f52c85dd88db11d5e03e146a3a6"},
+ {file = "mypy-0.812-cp36-cp36m-win_amd64.whl", hash = "sha256:25adde9b862f8f9aac9d2d11971f226bd4c8fbaa89fb76bdadb267ef22d10064"},
+ {file = "mypy-0.812-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:552a815579aa1e995f39fd05dde6cd378e191b063f031f2acfe73ce9fb7f9e56"},
+ {file = "mypy-0.812-cp37-cp37m-manylinux1_x86_64.whl", hash = "sha256:499c798053cdebcaa916eef8cd733e5584b5909f789de856b482cd7d069bdad8"},
+ {file = "mypy-0.812-cp37-cp37m-manylinux2010_x86_64.whl", hash = "sha256:5873888fff1c7cf5b71efbe80e0e73153fe9212fafdf8e44adfe4c20ec9f82d7"},
+ {file = "mypy-0.812-cp37-cp37m-win_amd64.whl", hash = "sha256:9f94aac67a2045ec719ffe6111df543bac7874cee01f41928f6969756e030564"},
+ {file = "mypy-0.812-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:d23e0ea196702d918b60c8288561e722bf437d82cb7ef2edcd98cfa38905d506"},
+ {file = "mypy-0.812-cp38-cp38-manylinux1_x86_64.whl", hash = "sha256:674e822aa665b9fd75130c6c5f5ed9564a38c6cea6a6432ce47eafb68ee578c5"},
+ {file = "mypy-0.812-cp38-cp38-manylinux2010_x86_64.whl", hash = "sha256:abf7e0c3cf117c44d9285cc6128856106183938c68fd4944763003decdcfeb66"},
+ {file = "mypy-0.812-cp38-cp38-win_amd64.whl", hash = "sha256:0d0a87c0e7e3a9becdfbe936c981d32e5ee0ccda3e0f07e1ef2c3d1a817cf73e"},
+ {file = "mypy-0.812-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:7ce3175801d0ae5fdfa79b4f0cfed08807af4d075b402b7e294e6aa72af9aa2a"},
+ {file = "mypy-0.812-cp39-cp39-manylinux1_x86_64.whl", hash = "sha256:b09669bcda124e83708f34a94606e01b614fa71931d356c1f1a5297ba11f110a"},
+ {file = "mypy-0.812-cp39-cp39-manylinux2010_x86_64.whl", hash = "sha256:33f159443db0829d16f0a8d83d94df3109bb6dd801975fe86bacb9bf71628e97"},
+ {file = "mypy-0.812-cp39-cp39-win_amd64.whl", hash = "sha256:3f2aca7f68580dc2508289c729bd49ee929a436208d2b2b6aab15745a70a57df"},
+ {file = "mypy-0.812-py3-none-any.whl", hash = "sha256:2f9b3407c58347a452fc0736861593e105139b905cca7d097e413453a1d650b4"},
+ {file = "mypy-0.812.tar.gz", hash = "sha256:cd07039aa5df222037005b08fbbfd69b3ab0b0bd7a07d7906de75ae52c4e3119"},
]
mypy-extensions = [
{file = "mypy_extensions-0.4.3-py2.py3-none-any.whl", hash = "sha256:090fedd75945a69ae91ce1303b5824f428daf5a028d2f6ab8a299250a846f15d"},
@@ -1543,8 +1548,8 @@ pygments = [
{file = "Pygments-2.8.0.tar.gz", hash = "sha256:37a13ba168a02ac54cc5891a42b1caec333e59b66addb7fa633ea8a6d73445c0"},
]
pylint = [
- {file = "pylint-2.6.2-py3-none-any.whl", hash = "sha256:e71c2e9614a4f06e36498f310027942b0f4f2fde20aebb01655b31edc63b9eaf"},
- {file = "pylint-2.6.2.tar.gz", hash = "sha256:718b74786ea7ed07aa0c58bf572154d4679f960d26e9641cc1de204a30b87fc9"},
+ {file = "pylint-2.7.1-py3-none-any.whl", hash = "sha256:a251b238db462b71d25948f940568bb5b3ae0e37dbaa05e10523f54f83e6cc7e"},
+ {file = "pylint-2.7.1.tar.gz", hash = "sha256:81ce108f6342421169ea039ff1f528208c99d2e5a9c4ca95cfc5291be6dfd982"},
]
pymdown-extensions = [
{file = "pymdown-extensions-7.1.tar.gz", hash = "sha256:5bf93d1ccd8281948cd7c559eb363e59b179b5373478e8a7195cf4b78e3c11b6"},
@@ -1681,8 +1686,8 @@ smmap = [
{file = "smmap-3.0.5.tar.gz", hash = "sha256:84c2751ef3072d4f6b2785ec7ee40244c6f45eb934d9e543e2c51f1bd3d54c50"},
]
tabulate = [
- {file = "tabulate-0.8.8-py3-none-any.whl", hash = "sha256:d6fe298fc0a58d848d6160118d17e70905f36766552ee78f8a1f4d64e8e16916"},
- {file = "tabulate-0.8.8.tar.gz", hash = "sha256:26f2589d80d332fefd2371d396863dedeb806f51b54bdb4b264579270b621e92"},
+ {file = "tabulate-0.8.9-py3-none-any.whl", hash = "sha256:d7c013fe7abbc5e491394e10fa845f8f32fe54f8dc60c6622c6cf482d25d47e4"},
+ {file = "tabulate-0.8.9.tar.gz", hash = "sha256:eb1d13f25760052e8931f2ef80aaf6045a6cceb47514db8beab24cded16f13a7"},
]
termcolor = [
{file = "termcolor-1.1.0.tar.gz", hash = "sha256:1d6d69ce66211143803fbc56652b41d73b4a400a2891d7bf7a1cdf4c02de613b"},
diff --git a/pyhdtoolkit/__init__.py b/pyhdtoolkit/__init__.py
index bb5074e6..c9536b36 100644
--- a/pyhdtoolkit/__init__.py
+++ b/pyhdtoolkit/__init__.py
@@ -13,7 +13,7 @@
__title__ = "pyhdtoolkit"
__description__ = "An all-in-one toolkit package to easy my Python work in my PhD."
__url__ = "https://github.com/fsoubelet/PyhDToolkit"
-__version__ = "0.8.1"
+__version__ = "0.8.2"
__author__ = "Felix Soubelet"
__author_email__ = "felix.soubelet@cern.ch"
__license__ = "MIT"
diff --git a/pyhdtoolkit/cpymadtools/latwiss.py b/pyhdtoolkit/cpymadtools/latwiss.py
index 89a12614..f1243417 100644
--- a/pyhdtoolkit/cpymadtools/latwiss.py
+++ b/pyhdtoolkit/cpymadtools/latwiss.py
@@ -18,7 +18,7 @@
from cpymad.madx import Madx
from loguru import logger
-from pyhdtoolkit.plotting.settings import PLOT_PARAMS
+from pyhdtoolkit.utils.defaults import PLOT_PARAMS
plt.rcParams.update(PLOT_PARAMS)
@@ -33,7 +33,7 @@ def _plot_lattice_series(
v_offset: float = 0.0,
color: str = "r",
alpha: float = 0.5,
- lw: int = 3,
+ **kwargs,
) -> None:
"""
Plots the layout of your machine as a patch of rectangles for different element types.
@@ -46,16 +46,19 @@ def _plot_lattice_series(
v_offset (float): vertical offset for the patch.
color (str): color kwarg to transmit to pyplot.
alpha (float): alpha kwarg to transmit to pyplot.
- lw (int): linewidth kwarg to transmit to pyplot.
+
+ Keyword Args:
+ Any kwarg that can be given to matplotlib.patches.Rectangle(), for instance `lw` for the edge line
+ width.
"""
ax.add_patch(
patches.Rectangle(
- (series.s - series.l, v_offset - height / 2.0),
+ (series.s - series.l, v_offset - height / 2.0), # anchor point
series.l, # width
height, # height
color=color,
alpha=alpha,
- lw=lw,
+ **kwargs,
)
)
@@ -63,25 +66,26 @@ def _plot_lattice_series(
def plot_latwiss(
madx: Madx,
title: str,
- figsize: Tuple[int, int] = (16, 10),
+ figsize: Tuple[int, int] = (18, 11),
savefig: str = None,
xlimits: Tuple[float, float] = None,
plot_dipoles: bool = True,
plot_quadrupoles: bool = True,
- plot_sextupoles: bool = False,
+ plot_bpms: bool = False,
disp_ylim: Tuple[float, float] = (-10, 125),
beta_ylim: Tuple[float, float] = None,
k0l_lim: Tuple[float, float] = (-0.25, 0.25),
k1l_lim: Tuple[float, float] = (-0.08, 0.08),
+ k2l_lim: Tuple[float, float] = None,
+ **kwargs,
) -> matplotlib.figure.Figure:
"""
- Provided with an active Cpymad class after having ran a script, will create a plot
- representing nicely the lattice layout and the beta functions along with the horizontal
- dispertion function. This is heavily refactored code, but the original is from Guido
- Sterbini.
+ Provided with an active Cpymad class after having ran a script, will create a plot representing nicely
+ the lattice layout and the beta functions along with the horizontal dispertion function. This is very
+ heavily refactored code, inspired by code from Guido Sterbini.
- WARNING: This will FAIL if you have not included 'q' or 'Q' in your quadrupoles' names,
- and 'b' or 'B' in your dipoles' names when defining your MAD-X sequence.
+ WARNING: This WILL FAIL if you have not included 'q' or 'Q' in your quadrupoles' names, and 'b' or 'B'
+ in your dipoles' names when defining your MAD-X sequence.
Args:
madx (cpymad.madx.Madx): an instanciated cpymad Madx object.
@@ -94,8 +98,8 @@ def plot_latwiss(
the figure. Defaults to True. Dipoles are plotted in blue.
plot_quadrupoles (bool): if True, quadrupole patches will be plotted on the layout
subplot of the figure. Defaults to True. Quadrupoles are plotted in red.
- plot_sextupoles (bool): if True, sextupole patches will be plotted on the layout subplot
- of the figure. Defaults to False. Sextupoles are plotted in yellow.
+ plot_bpms (bool): if True, additional patches will be plotted on the layout subplot to represent
+ Beam Position Monitors. BPMs are plotted in dark grey.
disp_ylim (Tuple[float, float]): vertical axis limits for the dispersion values.
Defaults to (-10, 125).
beta_ylim (Tuple[float, float]): vertical axis limits for the betatron function values.
@@ -104,72 +108,136 @@ def plot_latwiss(
height of dipole patches. Defaults to (-0.25, 0.25).
k1l_lim (Tuple[float, float]): vertical axis limits for the k1l values used for the
height of quadrupole patches. Defaults to (-0.08, 0.08).
+ k2l_lim (Tuple[float, float]): if given, sextupole patches will be plotted on the layout subplot of
+ the figure, and the provided values act as vertical axis limits for the k2l values used for the
+ height of sextupole patches.
+
+ Keyword Args:
+ Any keyword argument to be transmitted to `_plot_lattice_series`, and later on to
+ `matplotlib.patches.Rectangle`, such as lw etc.
+
+ WARNING:
+ Currently the function tries to plot legends for the different layout patches. The position of the
+ different legends has been hardcoded and can lead to messed-up layouts. User beware.
Returns:
The figure on which the plots are drawn. The underlying axes can be accessed with
'fig.get_axes()'. Eventually saves the figure as a file.
"""
# pylint: disable=too-many-arguments
+ # Restrict the span of twiss_df to avoid plotting all elements then cropping when xlimits is given
+ logger.info("Plotting optics functions and machine layout")
logger.debug("Getting Twiss dataframe from cpymad")
- twiss_df = madx.table.twiss.dframe()
+ twiss_df = madx.table.twiss.dframe().copy()
+ twiss_df = twiss_df[twiss_df.s.between(xlimits[0], xlimits[1])] if xlimits else twiss_df
figure = plt.figure(figsize=figsize)
# Create a subplot for the lattice patches (takes a third of figure)
logger.trace("Setting up element patches subplots")
quadrupole_patches_axis = plt.subplot2grid((3, 3), (0, 0), colspan=3, rowspan=1)
- dipole_patches_axis = quadrupole_patches_axis.twinx()
quadrupole_patches_axis.set_ylabel("1/f=K1L [m$^{-1}$]", color="red") # quadrupole in red
quadrupole_patches_axis.tick_params(axis="y", labelcolor="red")
- dipole_patches_axis.set_ylabel("$\\theta$=K0L [rad]", color="blue") # dipoles in blue
- dipole_patches_axis.tick_params(axis="y", labelcolor="blue")
quadrupole_patches_axis.set_ylim(k1l_lim)
+ quadrupole_patches_axis.set_title(title)
+ quadrupole_patches_axis.plot(twiss_df.s, 0 * twiss_df.s, "k") # 0-level line
+
+ dipole_patches_axis = quadrupole_patches_axis.twinx()
+ dipole_patches_axis.set_ylabel("$\\theta$=K0L [rad]", color="royalblue") # dipoles in blue
+ dipole_patches_axis.tick_params(axis="y", labelcolor="royalblue")
dipole_patches_axis.set_ylim(k0l_lim)
dipole_patches_axis.grid(False)
- quadrupole_patches_axis.set_title(title)
- quadrupole_patches_axis.plot(twiss_df.s, 0 * twiss_df.s, "k")
- # All elements can be defined as a 'multipole', but dipoles can also be defined as
- # 'rbend' or 'sbend', quadrupoles as 'quadrupoles' and sextupoles as 'sextupoles'
+ # All elements can be defined as a 'multipole', but dipoles can also be defined as 'rbend' or 'sbend',
+ # quadrupoles as 'quadrupoles' and sextupoles as 'sextupoles'. Function does not handle higher orders.
logger.debug("Extracting element-specific dataframes")
- quadrupoles_df = twiss_df[
- (twiss_df.keyword.isin(["multipole", "quadrupole"])) & (twiss_df.name.str.contains("Q", case=False))
- ]
dipoles_df = twiss_df[
(twiss_df.keyword.isin(["multipole", "rbend", "sbend"]))
& (twiss_df.name.str.contains("B", case=False))
]
+ quadrupoles_df = twiss_df[
+ (twiss_df.keyword.isin(["multipole", "quadrupole"])) & (twiss_df.name.str.contains("Q", case=False))
+ ]
sextupoles_df = twiss_df[
(twiss_df.keyword.isin(["multipole", "sextupole"])) & (twiss_df.name.str.contains("S", case=False))
]
+ bpms_df = twiss_df[(twiss_df.keyword.isin(["monitor"])) & (twiss_df.name.str.contains("BPM", case=False))]
- # Plotting dipole patches, beware 'sbend' and 'rbend' have an 'angle' value and not a 'k0l'
- if plot_dipoles:
+ if plot_dipoles: # beware 'sbend' and 'rbend' have an 'angle' value and not a 'k0l'
logger.debug("Plotting dipole patches")
- for _, dipole in dipoles_df.iterrows():
- if dipole.k0l != 0:
- logger.trace("Plotting dipole element")
+ plotted_elements = 0 # will help us not declare a label for legend at every patch
+ for dipole_name, dipole in dipoles_df.iterrows(): # by default twiss.dframe() has names as index
+ if dipole.k0l != 0 or dipole.angle != 0:
+ logger.trace(f"Plotting dipole element '{dipole_name}'")
_plot_lattice_series(
- dipole_patches_axis, dipole, height=dipole.k0l, v_offset=dipole.k0l / 2, color="b",
- )
- if dipole.angle != 0:
- logger.trace("Plotting 'sbend' / 'rbend' element")
- _plot_lattice_series(
- dipole_patches_axis, dipole, height=dipole.angle, v_offset=dipole.angle / 2, color="b",
+ dipole_patches_axis,
+ dipole,
+ height=dipole.k0l if dipole.k0l != 0 else dipole.angle,
+ v_offset=dipole.k0l / 2 if dipole.k0l != 0 else dipole.angle / 2,
+ color="royalblue",
+ label="MB" if plotted_elements == 0 else None,
+ **kwargs,
)
+ plotted_elements += 1
+ dipole_patches_axis.legend(loc=1, fontsize=16)
- # Plotting the quadrupole patches
if plot_quadrupoles:
logger.debug("Plotting quadrupole patches")
- for _, quad in quadrupoles_df.iterrows():
+ plotted_elements = 0
+ for quadrupole_name, quadrupole in quadrupoles_df.iterrows():
+ logger.trace(f"Plotting quadrupole element '{quadrupole_name}'")
_plot_lattice_series(
- quadrupole_patches_axis, quad, height=quad.k1l, v_offset=quad.k1l / 2, color="r"
+ quadrupole_patches_axis,
+ quadrupole,
+ height=quadrupole.k1l,
+ v_offset=quadrupole.k1l / 2,
+ color="r",
+ label="MQ" if plotted_elements == 0 else None,
+ **kwargs,
)
+ plotted_elements += 1
+ quadrupole_patches_axis.legend(loc=2, fontsize=16)
- # Plotting the sextupole patches
- if plot_sextupoles:
+ if k2l_lim:
logger.debug("Plotting sextupole patches")
- for _, sext in sextupoles_df.iterrows():
- _plot_lattice_series(dipole_patches_axis, sext, height=sext.k2l, v_offset=sext.k2l / 2, color="y")
+ sextupoles_patches_axis = quadrupole_patches_axis.twinx()
+ sextupoles_patches_axis.set_ylabel("K2L [m$^{-2}$]", color="darkgoldenrod")
+ sextupoles_patches_axis.tick_params(axis="y", labelcolor="darkgoldenrod")
+ sextupoles_patches_axis.spines["right"].set_position(("axes", 1.1))
+ sextupoles_patches_axis.set_ylim(k2l_lim)
+ plotted_elements = 0
+ for sextupole_name, sextupole in sextupoles_df.iterrows():
+ logger.trace(f"Plotting sextupole element '{sextupole_name}'")
+ _plot_lattice_series(
+ sextupoles_patches_axis,
+ sextupole,
+ height=sextupole.k2l,
+ v_offset=sextupole.k2l / 2,
+ color="goldenrod",
+ label="MS" if plotted_elements == 0 else None,
+ **kwargs,
+ )
+ plotted_elements += 1
+ sextupoles_patches_axis.legend(loc=3, fontsize=16)
+
+ if plot_bpms:
+ logger.debug("Plotting BPM patches")
+ bpm_patches_axis = quadrupole_patches_axis.twinx()
+ bpm_patches_axis.set_axis_off() # hide yticks, labels etc
+ bpm_patches_axis.set_ylim(-1.6, 1.6)
+ plotted_elements = 0
+ for bpm_name, bpm in bpms_df.iterrows():
+ logger.trace(f"Plotting BPM element '{bpm_name}'")
+ _plot_lattice_series(
+ bpm_patches_axis,
+ bpm,
+ height=2,
+ v_offset=0,
+ color="dimgrey",
+ label="BPM" if plotted_elements == 0 else None,
+ **kwargs,
+ )
+ plotted_elements += 1
+ bpm_patches_axis.legend(loc=4, fontsize=16)
# Plotting beta functions on remaining two thirds of the figure
logger.trace("Setting up betatron functions subplot")
@@ -178,13 +246,8 @@ def plot_latwiss(
betatron_axis.plot(twiss_df.s, twiss_df.bety, label="$\\beta_y$", lw=1.5)
betatron_axis.legend(loc=2)
betatron_axis.set_ylabel("$\\beta$-functions [m]")
-
- if beta_ylim:
- logger.debug("Setting ylim for betatron functions plot")
- betatron_axis.set_ylim(beta_ylim)
betatron_axis.set_xlabel("s [m]")
- # Plotting the dispersion
logger.trace("Setting up dispersion functions subplot")
dispertion_axis = betatron_axis.twinx()
dispertion_axis.plot(twiss_df.s, twiss_df.dx, color="brown", label="$D_x$", lw=2)
@@ -194,6 +257,10 @@ def plot_latwiss(
dispertion_axis.tick_params(axis="y", labelcolor="brown")
dispertion_axis.grid(False)
+ if beta_ylim:
+ logger.debug("Setting ylim for betatron functions plot")
+ betatron_axis.set_ylim(beta_ylim)
+
if disp_ylim:
logger.debug("Setting ylim for dispersion plot")
dispertion_axis.set_ylim(disp_ylim)
@@ -204,7 +271,7 @@ def plot_latwiss(
if savefig:
logger.info(f"Saving latwiss plot as {savefig}")
- plt.savefig(savefig, format="pdf", dpi=500)
+ plt.savefig(savefig)
return figure
@@ -235,7 +302,7 @@ def plot_machine_survey(
'fig.get_axes()'. Eventually saves the figure as a file.
"""
logger.debug("Getting machine survey from cpymad")
- madx.input("survey;")
+ madx.command.survey()
survey = madx.table.survey.dframe()
figure = plt.figure(figsize=figsize)
diff --git a/pyhdtoolkit/cpymadtools/optics.py b/pyhdtoolkit/cpymadtools/optics.py
deleted file mode 100644
index 5fad3939..00000000
--- a/pyhdtoolkit/cpymadtools/optics.py
+++ /dev/null
@@ -1,16 +0,0 @@
-"""
-Module cpymadtools.optics
--------------------------
-
-Created on 2020.02.19
-:author: Felix Soubelet (felix.soubelet@cern.ch)
-
-A module with functions to access `optics_functions` functionality directly onto `cpymad.madx.Madx` objects.
-"""
-from typing import Sequence
-
-import tfs
-import optics_functions
-
-from cpymad.madx import Madx
-from loguru import logger
diff --git a/pyhdtoolkit/cpymadtools/plotters.py b/pyhdtoolkit/cpymadtools/plotters.py
index 8a9c6dc2..c1f4cd8e 100644
--- a/pyhdtoolkit/cpymadtools/plotters.py
+++ b/pyhdtoolkit/cpymadtools/plotters.py
@@ -21,7 +21,7 @@
from matplotlib import colors as mcolors
from pyhdtoolkit.optics.twiss import courant_snyder_transform
-from pyhdtoolkit.plotting.settings import PLOT_PARAMS
+from pyhdtoolkit.utils.defaults import PLOT_PARAMS
plt.rcParams.update(PLOT_PARAMS)
diff --git a/pyhdtoolkit/cpymadtools/special.py b/pyhdtoolkit/cpymadtools/special.py
index fd9f5bbd..7598b4f0 100644
--- a/pyhdtoolkit/cpymadtools/special.py
+++ b/pyhdtoolkit/cpymadtools/special.py
@@ -139,6 +139,11 @@ def apply_lhc_rigidity_waist_shift_knob(
Applies the LHC rigidity waist shift knob, moving the waist left or right of IP. If you don't know what
this is, you should not be using this function.
+ Warning: Applying the shift will modify your tunes and most likely flip them, making a subsequent
+ matching impossible if your lattice has coupling. To avoid this, match to tunes split further apart
+ before applying the waist shift knob, and then match to the desired working point. For instance for
+ the LHC, match to (62.27, 60.36) before applying, and only then match to (62.31, 60.32).
+
Args:
madx (cpymad.madx.Madx): an instanciated cpymad Madx object.
rigidty_waist_shift_value (float): Units of the rigidity waist shift knob (positive values only).
diff --git a/pyhdtoolkit/cpymadtools/twiss.py b/pyhdtoolkit/cpymadtools/twiss.py
index 403cd2a8..f067900a 100644
--- a/pyhdtoolkit/cpymadtools/twiss.py
+++ b/pyhdtoolkit/cpymadtools/twiss.py
@@ -21,7 +21,7 @@
def get_pattern_twiss(
- madx: Madx, patterns: Sequence[str], columns: Sequence[str] = DEFAULT_TWISS_COLUMNS, **kwargs
+ madx: Madx, patterns: Sequence[str] = [""], columns: Sequence[str] = None, **kwargs,
) -> tfs.TfsDataFrame:
"""
Extract the `TWISS` table for desired variables, and for certain elements matching a pattern.
@@ -35,8 +35,10 @@ def get_pattern_twiss(
Args:
madx (cpymad.madx.Madx): an instanciated cpymad Madx object.
patterns (Sequence[str]): the different element patterns (such as `MQX` or `BPM`) to be applied to
- the command, which will determine the rows in the returned DataFrame.
- columns (Sequence[str]): the variables to be returned, as columns in the DataFrame.
+ the command, which will determine the rows in the returned DataFrame. Defaults to [""] which
+ will select all elements.
+ columns (Sequence[str]): the variables to be returned, as columns in the DataFrame. Defaults to
+ None, which will return all available columns.
Keyword Args:
Any keyword argument that can be given to the MAD-X TWISS command, such as `chrom`, `ripken`,
@@ -55,7 +57,7 @@ def get_pattern_twiss(
logger.trace("Extracting relevant parts of the TWISS table")
twiss_df = tfs.TfsDataFrame(madx.table.twiss.dframe().copy())
- twiss_df.headers = {var: madx.table.summ[var][0] for var in madx.table.summ}
+ twiss_df.headers = {var.upper(): madx.table.summ[var][0] for var in madx.table.summ}
twiss_df = twiss_df[madx.table.twiss.selected_columns()].iloc[
np.array(madx.table.twiss.selected_rows()).astype(bool)
]
diff --git a/pyhdtoolkit/maths/stats_fitting.py b/pyhdtoolkit/maths/stats_fitting.py
index 1fa1d867..e42f2b6d 100644
--- a/pyhdtoolkit/maths/stats_fitting.py
+++ b/pyhdtoolkit/maths/stats_fitting.py
@@ -13,6 +13,7 @@
from typing import Dict, Tuple, Union
import matplotlib
+import matplotlib.pyplot as plt # if omitted, get AttributeError: module 'matplotlib' has no attribute 'axes'
import numpy as np
import pandas as pd
import scipy.stats as st
diff --git a/pyhdtoolkit/maths/utils.py b/pyhdtoolkit/maths/utils.py
new file mode 100644
index 00000000..694e080f
--- /dev/null
+++ b/pyhdtoolkit/maths/utils.py
@@ -0,0 +1,39 @@
+from typing import Tuple, Union
+
+import numpy as np
+import pandas as pd
+
+from loguru import logger
+
+# ----- Miscellaneous Utilites ----- #
+
+
+def get_magnitude(value: float) -> int:
+ """Return the determined magnitude of the provided value."""
+ return int(np.floor(np.log10(np.abs(value))))
+
+
+def get_scaled_values_and_magnitude_string(
+ values_array: Union[pd.DataFrame, np.ndarray], force_magnitude: float = None
+) -> Tuple[Union[pd.DataFrame, np.ndarray], str]:
+ """
+ Conveniently scale provided values to the best determined magnitude. Returns scaled values
+ and the magnitude string to use in plots labels.
+
+ Args:
+ values_array (Union[pd.DataFrame, np.ndarray]): vectorised structure containing the
+ values to scale.
+ force_magnitude (float0: a specific magnitude value to use for the scaling, if desired.
+
+ Returns:
+ A tuple of the scaled values (same type as the provided ones) and the string to use for
+ the scale in plots labels and legends.
+
+ Usage:
+ """
+ magnitude = get_magnitude(max(values_array)) if force_magnitude is None else force_magnitude
+ applied_magnitude = -magnitude
+ logger.trace(f"Scaling data by {applied_magnitude} orders of magnitude")
+ scaled_values = values_array * (10 ** applied_magnitude)
+ magnitude_string = "{" + f"{applied_magnitude}" + "}"
+ return scaled_values, magnitude_string
diff --git a/pyhdtoolkit/optics/ripken.py b/pyhdtoolkit/optics/ripken.py
new file mode 100644
index 00000000..ce49e088
--- /dev/null
+++ b/pyhdtoolkit/optics/ripken.py
@@ -0,0 +1,53 @@
+from typing import Union
+
+import numba
+import numpy as np
+
+from loguru import logger
+
+# ----- Setup Utilites ----- #
+
+
+def lebedev_beam_size(
+ beta1_: Union[float, np.ndarray], beta2_: Union[float, np.ndarray], geom_emit_x: float, geom_emit_y: float
+) -> Union[float, np.ndarray]:
+ """
+ Calculate beam size according to the Lebedev-bogacz formula, based on the Ripken-Mais
+ Twiss parameters. The implementation is that of Eq. (A.3.1) in FERMILAB-PUB-10-383-AD, avaliable at the
+ following link: https://arxiv.org/ftp/arxiv/papers/1207/1207.5526.pdf
+
+ Args:
+ beta1_ (Union[float, np.ndarray]): value(s) for the beta1x or beta1y Ripken parameter.
+ beta2_ (Union[float, np.ndarray]): value(s) for the beta2x or beta2y Ripken parameter.
+ geom_emit_x (float): geometric emittance of the horizontal plane.
+ geom_emit_y (float): geometric emittante of the vertical plane.
+
+ Returns:
+ The beam size (horizontal or vertical) according to Lebedev & Bogacz, as sqrt(epsx *
+ beta1_^2 + epsy * beta2_^2).
+ """
+ logger.trace("Computing beam size according to Lebedev formula: sqrt(epsx * b1_^2 + epsy * b2_^2)")
+ return np.sqrt(geom_emit_x * beta1_ + geom_emit_y * beta2_)
+
+
+# ----- JITed Calculations ----- #
+
+
+@numba.njit()
+def _beam_size(coordinates_distribution: np.ndarray, method: str = "std") -> float:
+ """
+ Compute beam size from particle coordinates.
+
+ Args:
+ coordinates_distribution (np.ndarray): ensemble of coordinates of the particle distributon.
+ method (str): the method of calculation to use, either 'std' (using the standard deviation as the
+ beam size) or 'rms' (root mean square).
+
+ Returns:
+ The computed beam size.
+ """
+ if method == "std":
+ return coordinates_distribution.std()
+ elif method == "rms":
+ return np.sqrt(np.mean(np.square(coordinates_distribution)))
+ raise NotImplementedError(f"Invalid method provided")
diff --git a/pyhdtoolkit/plotting/__init__.py b/pyhdtoolkit/plotting/__init__.py
index 46b93c8b..9fd6e57d 100644
--- a/pyhdtoolkit/plotting/__init__.py
+++ b/pyhdtoolkit/plotting/__init__.py
@@ -8,4 +8,3 @@
"""
from .helpers import AnnotationsPlotter
-from .settings import PLOT_PARAMS
diff --git a/pyhdtoolkit/plotting/helpers.py b/pyhdtoolkit/plotting/helpers.py
index 9eb6682a..d5f40f4e 100644
--- a/pyhdtoolkit/plotting/helpers.py
+++ b/pyhdtoolkit/plotting/helpers.py
@@ -26,6 +26,7 @@ def set_arrow_label(
color: str = "k",
arrow_arc_rad: float = -0.2,
fontsize: int = 20,
+ **kwargs,
) -> matplotlib.text.Annotation:
"""
Add a label box with text and an arrow from the box to a specified position to an existing
@@ -59,4 +60,5 @@ def set_arrow_label(
arrowprops=dict(
arrowstyle="-|>", connectionstyle="arc3,rad=" + str(arrow_arc_rad), fc="w", color=color, lw=2,
),
+ **kwargs,
)
diff --git a/pyhdtoolkit/plotting/settings.py b/pyhdtoolkit/plotting/settings.py
deleted file mode 100644
index f49b9d81..00000000
--- a/pyhdtoolkit/plotting/settings.py
+++ /dev/null
@@ -1,63 +0,0 @@
-"""
-Module plotting.settings
-------------------------
-
-Created on 2019.12.08
-:author: Felix Soubelet (felix.soubelet@cern.ch)
-
-Some settings for better matplotlib.pyplot plots.
-Work in progress.
-"""
-from typing import Dict, Union
-
-# Set those with matplotlib.pyplot.rcParams.update(PLOT_PARAMS).
-# Will ALWAYS be overwritten by later on definition
-PLOT_PARAMS: Dict[str, Union[float, bool, str, tuple]] = {
- # ------ Axes ------ #
- "axes.linewidth": 0.8, # Linewidth of axes edges
- "axes.grid": False, # Do not display grid
- "axes.labelsize": 25, # Fontsize of the x and y axis labels
- "axes.titlesize": 27, # Fontsize of the axes title
- # ------ Date Forrmats ------ #
- "date.autoformatter.year": "%Y", # AutoDateFormatter setting for years display
- "date.autoformatter.month": "%Y-%m", # AutoDateFormatter setting for months display
- "date.autoformatter.day": "%Y-%m-%d", # AutoDateFormatter setting for days display
- "date.autoformatter.hour": "%m-%d %H", # AutoDateFormatter setting for seconds display
- "date.autoformatter.minute": "%d %H:%M", # AutoDateFormatter setting for minutes display
- "date.autoformatter.second": "%H:%M:%S", # AutoDateFormatter setting for seconds display
- "date.autoformatter.microsecond": "%M:%S.%f", # AutoDateFormatter setting for microseconds
- # ------ General Figure ------ #
- "figure.autolayout": True, # Adjust subplot params to fit the figure (tight_layout)
- "figure.dpi": 300, # Figure dots per inch
- "figure.figsize": (18, 11), # Size of the figure
- "figure.max_open_warning": 10, # Max number of figures to open before warning
- "figure.titlesize": 30, # Size of the figure title
- # ------ Fonts ------ #
- "font.family": "sans-serif", # Font family
- # "font.sans-serif": "Helvetica", # Sans-Serif font to use
- "font.style": "normal", # Style to apply to text font
- # ------- Legend ------ #
- "legend.fancybox": True, # Use rounded box for legend background
- "legend.fontsize": 22, # Legend text font size
- "legend.loc": "best", # Default legend location
- # ------ Lines ------ #
- "lines.linewidth": 1, # Line width, in points
- "lines.markersize": 5, # Marker size, in points
- "lines.antialiased": True, # Apply anti-aliasing to lines display
- # ------ Patches ------ #
- "patch.linewidth": 1, # Width of patches edge lines
- "patch.antialiased": True, # Apply anti-aliasing to patches display
- # ------ Paths ------ #
- "path.simplify": True, # Reduce file size by removing "invisible" points
- # ------ Saving ------ #
- "savefig.dpi": 300, # Saved figure dots per inch
- "savefig.format": "pdf", # Saved figure file format
- "savefig.bbox": "tight", # Careful: incompatible with pipe-based animation backends
- # ------ Text ------ #
- "text.antialiased": True, # Apply anti-aliasing to text elements
- "text.color": "black", # Default text color
- "text.usetex": False, # Do not use LaTeX for text handling (I don't have a local installation)
- # ------ Ticks ------ #
- "xtick.labelsize": 20, # Fontsize of the x axis tick labels
- "ytick.labelsize": 20, # Fontsize of the y axis tick labels
-}
diff --git a/pyhdtoolkit/tfstools/latwiss.py b/pyhdtoolkit/tfstools/latwiss.py
index 06892080..d0fce706 100644
--- a/pyhdtoolkit/tfstools/latwiss.py
+++ b/pyhdtoolkit/tfstools/latwiss.py
@@ -19,7 +19,7 @@
from loguru import logger
-from pyhdtoolkit.plotting.settings import PLOT_PARAMS
+from pyhdtoolkit.utils.defaults import PLOT_PARAMS
plt.rcParams.update(PLOT_PARAMS)
diff --git a/pyhdtoolkit/utils/defaults.py b/pyhdtoolkit/utils/defaults.py
index 2b3ec323..da5fb418 100644
--- a/pyhdtoolkit/utils/defaults.py
+++ b/pyhdtoolkit/utils/defaults.py
@@ -7,7 +7,12 @@
Provides defaults to import for different settings.
"""
+import sys
+
from pathlib import Path
+from typing import Dict, Union
+
+from loguru import logger
ANACONDA_INSTALL = Path().home() / "anaconda3"
OMC_PYTHON = ANACONDA_INSTALL / "envs" / "OMC" / "bin" / "python"
@@ -24,3 +29,67 @@
"{name}:{line} - "
"{message}"
)
+
+# Set those with matplotlib.pyplot.rcParams.update(PLOT_PARAMS).
+# Will ALWAYS be overwritten by later on definition
+PLOT_PARAMS: Dict[str, Union[float, bool, str, tuple]] = {
+ # ------ Axes ------ #
+ "axes.linewidth": 0.8, # Linewidth of axes edges
+ "axes.grid": False, # Do not display grid
+ "axes.labelsize": 25, # Fontsize of the x and y axis labels
+ "axes.titlesize": 27, # Fontsize of the axes title
+ "axes.formatter.limits": (-3, 5), # Switch to scientific notations when order of magnitude reaches 1e3
+ # "axes.formatter.useoffset": False, # Do not use the annoying offset on top of yticks
+ "axes.formatter.use_mathtext": True, # Format with i.e 10^{4} instead of 1e4
+ # ------ Date Formats ------ #
+ "date.autoformatter.year": "%Y", # AutoDateFormatter setting for years display
+ "date.autoformatter.month": "%Y-%m", # AutoDateFormatter setting for months display
+ "date.autoformatter.day": "%Y-%m-%d", # AutoDateFormatter setting for days display
+ "date.autoformatter.hour": "%m-%d %H", # AutoDateFormatter setting for seconds display
+ "date.autoformatter.minute": "%d %H:%M", # AutoDateFormatter setting for minutes display
+ "date.autoformatter.second": "%H:%M:%S", # AutoDateFormatter setting for seconds display
+ "date.autoformatter.microsecond": "%M:%S.%f", # AutoDateFormatter setting for microseconds
+ # ------ General Figure ------ #
+ "figure.autolayout": True, # Adjust subplot params to fit the figure (tight_layout)
+ "figure.dpi": 300, # Figure dots per inch
+ "figure.figsize": (18, 11), # Size of the figure
+ "figure.max_open_warning": 10, # Max number of figures to open before warning
+ "figure.titlesize": 30, # Size of the figure title
+ # ------ Fonts ------ #
+ "font.family": "sans-serif", # Font family
+ # "font.sans-serif": "Helvetica", # Sans-Serif font to use
+ "font.style": "normal", # Style to apply to text font
+ # ------- Legend ------ #
+ "legend.fancybox": True, # Use rounded box for legend background
+ "legend.fontsize": 22, # Legend text font size
+ "legend.loc": "best", # Default legend location
+ # ------ Lines ------ #
+ "lines.linewidth": 1, # Line width, in points
+ "lines.markersize": 5, # Marker size, in points
+ "lines.antialiased": True, # Apply anti-aliasing to lines display
+ # ------ Patches ------ #
+ "patch.linewidth": 1, # Width of patches edge lines
+ "patch.antialiased": True, # Apply anti-aliasing to patches display
+ # ------ Paths ------ #
+ "path.simplify": True, # Reduce file size by removing "invisible" points
+ # ------ Saving ------ #
+ "savefig.dpi": 350, # Saved figure dots per inch
+ "savefig.format": "pdf", # Saved figure file format
+ "savefig.bbox": "tight", # Careful: incompatible with pipe-based animation backends
+ # ------ Text ------ #
+ "text.antialiased": True, # Apply anti-aliasing to text elements
+ "text.color": "black", # Default text color
+ "text.usetex": False, # Do not use LaTeX for text handling (I don't have a local installation)
+ # ------ Ticks ------ #
+ "xtick.labelsize": 20, # Fontsize of the x axis tick labels
+ "ytick.labelsize": 20, # Fontsize of the y axis tick labels
+}
+
+
+def config_logger(level: str = "INFO") -> None:
+ """
+ Resets the logger object from loguru, with `sys.stdout` as a sink and the aforedefined format.
+ This comes down to personnal preference.
+ """
+ logger.remove()
+ logger.add(sys.stdout, format=LOGURU_FORMAT, level=level.upper())
diff --git a/pyproject.toml b/pyproject.toml
index ff983700..2a7ee546 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
[tool.poetry]
name = "pyhdtoolkit"
-version = "0.8.1"
+version = "0.8.2"
description = "An all-in-one toolkit package to easy my Python work in my PhD."
authors = ["Felix Soubelet "]
license = "MIT"
@@ -40,7 +40,7 @@ tfs-pandas = "^2.0"
loguru = "^0"
numba = ">=0.51,<1.0"
cpymad = "^1.6"
-optics-functions = "^0"
+#optics-functions = "^0"
# Defining optional dependencies for extras
portray = { version = "^1.4", optional = true }
diff --git a/tests/baseline/test_plot_horizontal_courant_snyder_phase_space.png b/tests/baseline/test_plot_horizontal_courant_snyder_phase_space.png
index bea2a86e..b433ec15 100644
Binary files a/tests/baseline/test_plot_horizontal_courant_snyder_phase_space.png and b/tests/baseline/test_plot_horizontal_courant_snyder_phase_space.png differ
diff --git a/tests/baseline/test_plot_latwiss.png b/tests/baseline/test_plot_latwiss.png
index 6e236c2d..282f8194 100644
Binary files a/tests/baseline/test_plot_latwiss.png and b/tests/baseline/test_plot_latwiss.png differ
diff --git a/tests/baseline/test_plot_vertical_courant_snyder_phase_space.png b/tests/baseline/test_plot_vertical_courant_snyder_phase_space.png
index 5f4312ac..30f42d25 100644
Binary files a/tests/baseline/test_plot_vertical_courant_snyder_phase_space.png and b/tests/baseline/test_plot_vertical_courant_snyder_phase_space.png differ
diff --git a/tests/baseline/test_set_arrow_label.png b/tests/baseline/test_set_arrow_label.png
deleted file mode 100644
index e559bba9..00000000
Binary files a/tests/baseline/test_set_arrow_label.png and /dev/null differ
diff --git a/tests/inputs/beta11.npy b/tests/inputs/beta11.npy
new file mode 100644
index 00000000..2cf37843
Binary files /dev/null and b/tests/inputs/beta11.npy differ
diff --git a/tests/inputs/beta21.npy b/tests/inputs/beta21.npy
new file mode 100644
index 00000000..59972008
Binary files /dev/null and b/tests/inputs/beta21.npy differ
diff --git a/tests/inputs/force_scaled.npy b/tests/inputs/force_scaled.npy
new file mode 100644
index 00000000..f69a6dfc
Binary files /dev/null and b/tests/inputs/force_scaled.npy differ
diff --git a/tests/inputs/ips_twiss.tfs b/tests/inputs/ips_twiss.tfs
index af058994..116a320d 100644
--- a/tests/inputs/ips_twiss.tfs
+++ b/tests/inputs/ips_twiss.tfs
@@ -1,30 +1,30 @@
-@ length %le 26658.88319999891
-@ orbit5 %le -0.0
-@ alfa %le 0.00034847119689069634
-@ gammatr %le 53.56937216865596
-@ q1 %le 62.30999999999979
-@ dq1 %le 1.9999511151009346
-@ betxmax %le 8012.910936548429
-@ dxmax %le 3.1878834016412942
-@ dxrms %le 1.5263586503155346
-@ xcomax %le 0.012145901206674433
-@ xcorms %le 0.0008469328809041056
-@ q2 %le 60.320000000000036
-@ dq2 %le 1.9989843817320194
-@ betymax %le 8012.925008813512
-@ dymax %le 3.0912647588025863
-@ dyrms %le 0.3523724306503348
-@ ycomax %le 0.012659398512083038
-@ ycorms %le 0.0009142417111966534
-@ deltap %le 0.0
-@ synch_1 %le 0.0
-@ synch_2 %le 0.0
-@ synch_3 %le 0.0
-@ synch_4 %le 0.0
-@ synch_5 %le 0.0
-@ synch_6 %le 0.0
-@ synch_8 %le 0.0
-@ nflips %le 0.0
+@ LENGTH %le 26658.88319999891
+@ ORBIT5 %le -0.0
+@ ALFA %le 0.00034847119689069634
+@ GAMMATR %le 53.56937216865596
+@ Q1 %le 62.30999999999979
+@ DQ1 %le 1.9999511151009346
+@ BETXMAX %le 8012.910936548429
+@ DXMAX %le 3.1878834016412942
+@ DXRMS %le 1.5263586503155346
+@ XCOMAX %le 0.012145901206674433
+@ XCORMS %le 0.0008469328809041056
+@ Q2 %le 60.320000000000036
+@ DQ2 %le 1.9989843817320194
+@ BETYMAX %le 8012.925008813512
+@ DYMAX %le 3.0912647588025863
+@ DYRMS %le 0.3523724306503348
+@ YCOMAX %le 0.012659398512083038
+@ YCORMS %le 0.0009142417111966534
+@ DELTAP %le 0.0
+@ SYNCH_1 %le 0.0
+@ SYNCH_2 %le 0.0
+@ SYNCH_3 %le 0.0
+@ SYNCH_4 %le 0.0
+@ SYNCH_5 %le 0.0
+@ SYNCH_6 %le 0.0
+@ SYNCH_8 %le 0.0
+@ NFLIPS %le 0.0
* name s x y px py betx bety alfx alfy dx dy mux muy r11 r12 r21 r22 beta11 beta12 beta21 beta22
$ %s %le %le %le %le %le %le %le %le %le %le %le %le %le %le %le %le %le %le %le %le %le
"ip1:1" 0 -0.000550000063398 2.11571696509e-12 -2.07589807937e-11 0.000160000110262 0.300000800432 0.300000217487 -1.39746377622e-06 6.301868104e-08 0.023825029733 0.00517636337348 0 0 1.42733657396e-05 -1.74256609927e-07 2.18087244679e-06 1.3724290696e-05 0 0 0 0
diff --git a/tests/inputs/ir1_twiss.tfs b/tests/inputs/ir1_twiss.tfs
index 11e7dbc6..a6c0861f 100644
--- a/tests/inputs/ir1_twiss.tfs
+++ b/tests/inputs/ir1_twiss.tfs
@@ -1,30 +1,30 @@
-@ length %le 26658.88319999891
-@ orbit5 %le -0.0
-@ alfa %le 0.00034847119689069634
-@ gammatr %le 53.56937216865596
-@ q1 %le 62.30999999999979
-@ dq1 %le 1.9999511151009346
-@ betxmax %le 8012.910936548429
-@ dxmax %le 3.1878834016412942
-@ dxrms %le 1.5263586503155346
-@ xcomax %le 0.012145901206674433
-@ xcorms %le 0.0008469328809041056
-@ q2 %le 60.320000000000036
-@ dq2 %le 1.9989843817320194
-@ betymax %le 8012.925008813512
-@ dymax %le 3.0912647588025863
-@ dyrms %le 0.3523724306503348
-@ ycomax %le 0.012659398512083038
-@ ycorms %le 0.0009142417111966534
-@ deltap %le 0.0
-@ synch_1 %le 0.0
-@ synch_2 %le 0.0
-@ synch_3 %le 0.0
-@ synch_4 %le 0.0
-@ synch_5 %le 0.0
-@ synch_6 %le 0.0
-@ synch_8 %le 0.0
-@ nflips %le 0.0
+@ LENGTH %le 26658.88319999891
+@ ORBIT5 %le -0.0
+@ ALFA %le 0.00034847119689069634
+@ GAMMATR %le 53.56937216865596
+@ Q1 %le 62.30999999999979
+@ DQ1 %le 1.9999511151009346
+@ BETXMAX %le 8012.910936548429
+@ DXMAX %le 3.1878834016412942
+@ DXRMS %le 1.5263586503155346
+@ XCOMAX %le 0.012145901206674433
+@ XCORMS %le 0.0008469328809041056
+@ Q2 %le 60.320000000000036
+@ DQ2 %le 1.9989843817320194
+@ BETYMAX %le 8012.925008813512
+@ DYMAX %le 3.0912647588025863
+@ DYRMS %le 0.3523724306503348
+@ YCOMAX %le 0.012659398512083038
+@ YCORMS %le 0.0009142417111966534
+@ DELTAP %le 0.0
+@ SYNCH_1 %le 0.0
+@ SYNCH_2 %le 0.0
+@ SYNCH_3 %le 0.0
+@ SYNCH_4 %le 0.0
+@ SYNCH_5 %le 0.0
+@ SYNCH_6 %le 0.0
+@ SYNCH_8 %le 0.0
+@ NFLIPS %le 0.0
* name s x y px py betx bety alfx alfy dx dy mux muy r11 r12 r21 r22 beta11 beta12 beta21 beta22
$ %s %le %le %le %le %le %le %le %le %le %le %le %le %le %le %le %le %le %le %le %le %le
"ip1:1" 0 -0.000550000063398 2.11571696509e-12 -2.07589807937e-11 0.000160000110262 0.300000800432 0.300000217487 -1.39746377622e-06 6.301868104e-08 0.023825029733 0.00517636337348 0 0 1.42733657396e-05 -1.74256609927e-07 2.18087244679e-06 1.3724290696e-05 0 0 0 0
diff --git a/tests/inputs/ir5_twiss.tfs b/tests/inputs/ir5_twiss.tfs
index 57776037..2b5bbc5a 100644
--- a/tests/inputs/ir5_twiss.tfs
+++ b/tests/inputs/ir5_twiss.tfs
@@ -1,30 +1,30 @@
-@ length %le 26658.88319999891
-@ orbit5 %le -0.0
-@ alfa %le 0.00034847119689069634
-@ gammatr %le 53.56937216865596
-@ q1 %le 62.30999999999979
-@ dq1 %le 1.9999511151009346
-@ betxmax %le 8012.910936548429
-@ dxmax %le 3.1878834016412942
-@ dxrms %le 1.5263586503155346
-@ xcomax %le 0.012145901206674433
-@ xcorms %le 0.0008469328809041056
-@ q2 %le 60.320000000000036
-@ dq2 %le 1.9989843817320194
-@ betymax %le 8012.925008813512
-@ dymax %le 3.0912647588025863
-@ dyrms %le 0.3523724306503348
-@ ycomax %le 0.012659398512083038
-@ ycorms %le 0.0009142417111966534
-@ deltap %le 0.0
-@ synch_1 %le 0.0
-@ synch_2 %le 0.0
-@ synch_3 %le 0.0
-@ synch_4 %le 0.0
-@ synch_5 %le 0.0
-@ synch_6 %le 0.0
-@ synch_8 %le 0.0
-@ nflips %le 0.0
+@ LENGTH %le 26658.88319999891
+@ ORBIT5 %le -0.0
+@ ALFA %le 0.00034847119689069634
+@ GAMMATR %le 53.56937216865596
+@ Q1 %le 62.30999999999979
+@ DQ1 %le 1.9999511151009346
+@ BETXMAX %le 8012.910936548429
+@ DXMAX %le 3.1878834016412942
+@ DXRMS %le 1.5263586503155346
+@ XCOMAX %le 0.012145901206674433
+@ XCORMS %le 0.0008469328809041056
+@ Q2 %le 60.320000000000036
+@ DQ2 %le 1.9989843817320194
+@ BETYMAX %le 8012.925008813512
+@ DYMAX %le 3.0912647588025863
+@ DYRMS %le 0.3523724306503348
+@ YCOMAX %le 0.012659398512083038
+@ YCORMS %le 0.0009142417111966534
+@ DELTAP %le 0.0
+@ SYNCH_1 %le 0.0
+@ SYNCH_2 %le 0.0
+@ SYNCH_3 %le 0.0
+@ SYNCH_4 %le 0.0
+@ SYNCH_5 %le 0.0
+@ SYNCH_6 %le 0.0
+@ SYNCH_8 %le 0.0
+@ NFLIPS %le 0.0
* name s x y px py betx bety alfx alfy dx dy mux muy r11 r12 r21 r22 beta11 beta12 beta21 beta22
$ %s %le %le %le %le %le %le %le %le %le %le %le %le %le %le %le %le %le %le %le %le %le
"mqxa.3l5:1" 13282.3242328 -0.00602728244789 -0.000311946997621 -0.000327181513215 -1.91355980396e-05 4925.77631446 5505.13023195 -256.399776417 301.496099112 -1.34949507914 2.554620677 30.7077802454 29.398936912 0.00350710250004 -0.0674939414892 -0.000191779310396 0.00369083631221 0 0 0 0
diff --git a/tests/inputs/lebedev_size.npy b/tests/inputs/lebedev_size.npy
new file mode 100644
index 00000000..04387b8a
Binary files /dev/null and b/tests/inputs/lebedev_size.npy differ
diff --git a/tests/inputs/scaled.npy b/tests/inputs/scaled.npy
new file mode 100644
index 00000000..2ba4dfe9
Binary files /dev/null and b/tests/inputs/scaled.npy differ
diff --git a/tests/inputs/to_scale.npy b/tests/inputs/to_scale.npy
new file mode 100644
index 00000000..b3f691fa
Binary files /dev/null and b/tests/inputs/to_scale.npy differ
diff --git a/tests/test_beam.py b/tests/test_beam.py
deleted file mode 100644
index 3c8a515a..00000000
--- a/tests/test_beam.py
+++ /dev/null
@@ -1,43 +0,0 @@
-import pytest
-
-from pyhdtoolkit.optics.beam import Beam
-
-
-class TestProperties:
- def test_gamma_rel(self):
- assert Beam(6500, 2.5e-6).gamma_rel == 6928.628011131436
-
- def test_beta_rel(self):
- assert Beam(6500, 2.5e-6).beta_rel == 1.0000000104153894
-
- def test_brho(self):
- assert Beam(6500, 2.5e-6).brho == 7.227222137900961e-05
-
- def test_normalized_emittance(self):
- assert Beam(6500, 2.5e-6).normalized_emittance == 0.01732157020823949
-
- def test_rms_emittance(self):
- assert Beam(6500, 2.5e-6).rms_emittance == 3.6082179183888383e-10
-
-
-class TestCalculations:
- def test_lhc_revolution_frequency(self):
- lhc_beam = Beam(6500, 2.5e-6)
- assert lhc_beam.revolution_frequency() == 11245.499628523643
-
- @pytest.mark.parametrize(
- "alpha_p, result",
- [(0, 2.083077890845299e-08), (1e-5, -9.979169221091548e-06), (-500, 500.0000000208308)],
- )
- def test_eta(self, alpha_p, result):
- assert Beam(6500, 2.5e-6).eta(alpha_p) == result
-
- @pytest.mark.parametrize(
- "alpha_p, result", [(1e-5, 316.2277660168379), (500, 0.044721359549995794)],
- )
- def test_gamma_transition(self, alpha_p, result):
- assert Beam(6500, 2.5e-6).gamma_transition(alpha_p) == result
-
- def test_gamma_transition_raises(self):
- with pytest.raises(ZeroDivisionError):
- Beam(6500, 2.5e-6).gamma_transition(0)
diff --git a/tests/test_cpymadtools.py b/tests/test_cpymadtools.py
index 776d8f6a..0e312045 100644
--- a/tests/test_cpymadtools.py
+++ b/tests/test_cpymadtools.py
@@ -52,6 +52,7 @@
make_lhc_thin,
make_sixtrack_output,
power_landau_octupoles,
+ re_cycle_sequence,
)
from pyhdtoolkit.cpymadtools.track import track_single_particle
@@ -211,7 +212,8 @@ def test_plot_latwiss(self, tmp_path):
title="Project 3 Base Lattice",
xlimits=(-50, 1_050),
beta_ylim=(5, 75),
- plot_sextupoles=True,
+ k2l_lim=(-0.25, 0.25),
+ plot_bpms=True,
savefig=saved_fig,
)
assert saved_fig.is_file()
@@ -772,6 +774,15 @@ def test_makethin_lhc(self, _matched_lhc_madx):
[coordinate in tracks.columns for coordinate in ("x", "px", "y", "py", "t", "pt", "s", "e")]
)
+ def test_re_cycling(self, _bare_lhc_madx):
+ madx = _bare_lhc_madx
+ re_cycle_sequence(madx, sequence="lhcb1", start="IP3")
+ make_lhc_beams(madx)
+ madx.command.use(sequence="lhcb1")
+ madx.twiss()
+ twiss = madx.table.twiss.dframe().copy()
+ assert "ip3" in twiss.name[0].lower()
+
class TestTrack:
def test_single_particle_tracking(self, _matched_base_lattice):
diff --git a/tests/test_maths.py b/tests/test_maths.py
index ac96f1cb..cfea27c6 100644
--- a/tests/test_maths.py
+++ b/tests/test_maths.py
@@ -1,3 +1,5 @@
+import pathlib
+
from copy import deepcopy
import numpy as np
@@ -7,7 +9,10 @@
import pyhdtoolkit.maths.nonconvex_phase_sync as nps
from pyhdtoolkit.maths import stats_fitting
+from pyhdtoolkit.maths import utils as mutils
+CURRENT_DIR = pathlib.Path(__file__).parent
+INPUTS_DIR = CURRENT_DIR / "inputs"
REF_DISTRIBUTIONS = deepcopy(stats_fitting.DISTRIBUTIONS)
@@ -133,6 +138,22 @@ def test_make_pdf(self, degrees_of_freedom):
pdf.idxmax() == pytest.approx(degrees_of_freedom - 2, rel=1e-2)
+class TestMagnitudeUtils:
+ @pytest.mark.parametrize("value, result", [(1, 0), (10, 1), (0.0311, -2), (5e-7, -7)])
+ def test_magnitude(self, value, result):
+ assert mutils.get_magnitude(value) == result
+
+ def test_mag_and_string(self, _to_scale, _scaled):
+ scaled, mag_str = mutils.get_scaled_values_and_magnitude_string(_to_scale)
+ assert np.allclose(scaled, _scaled)
+ assert mag_str == "{-1}"
+
+ def test_mag_and_string_forced_scale(self, _to_scale, _force_scaled):
+ scaled, mag_str = mutils.get_scaled_values_and_magnitude_string(_to_scale, force_magnitude=2)
+ assert np.allclose(scaled, _force_scaled)
+ assert mag_str == "{-2}"
+
+
# ---------------------- Utilities ---------------------- #
@@ -188,3 +209,18 @@ def _create_2d_gaussian_noise(mean: float, stdev: float, shape: tuple) -> np.nda
gaussian_2d_mat = np.random.default_rng().normal(mean, stdev, size=shape)
upper_triangle = np.triu(gaussian_2d_mat)
return upper_triangle - upper_triangle.T
+
+
+@pytest.fixture()
+def _to_scale() -> np.ndarray:
+ return np.load(INPUTS_DIR / "to_scale.npy")
+
+
+@pytest.fixture()
+def _scaled() -> np.ndarray:
+ return np.load(INPUTS_DIR / "scaled.npy")
+
+
+@pytest.fixture()
+def _force_scaled() -> np.ndarray:
+ return np.load(INPUTS_DIR / "force_scaled.npy")
diff --git a/tests/test_optics.py b/tests/test_optics.py
new file mode 100644
index 00000000..0bb00a88
--- /dev/null
+++ b/tests/test_optics.py
@@ -0,0 +1,107 @@
+import pathlib
+
+import numpy as np
+import pytest
+
+from pyhdtoolkit.optics import ripken, twiss
+from pyhdtoolkit.optics.beam import Beam
+
+CURRENT_DIR = pathlib.Path(__file__).parent
+INPUTS_DIR = CURRENT_DIR / "inputs"
+INPUT_PATHS = {
+ "alpha_beta": INPUTS_DIR / "alpha_beta.npy",
+ "u_vector": INPUTS_DIR / "u_vector.npy",
+ "u_bar": INPUTS_DIR / "u_bar.npy",
+ "beta11": INPUTS_DIR / "beta11.npy",
+ "beta21": INPUTS_DIR / "beta21.npy",
+ "lebedev": INPUTS_DIR / "lebedev_size.npy",
+}
+
+
+class TestBeamProperties:
+ def test_gamma_rel(self):
+ assert Beam(6500, 2.5e-6).gamma_rel == 6928.628011131436
+
+ def test_beta_rel(self):
+ assert Beam(6500, 2.5e-6).beta_rel == 1.0000000104153894
+
+ def test_brho(self):
+ assert Beam(6500, 2.5e-6).brho == 7.227222137900961e-05
+
+ def test_normalized_emittance(self):
+ assert Beam(6500, 2.5e-6).normalized_emittance == 0.01732157020823949
+
+ def test_rms_emittance(self):
+ assert Beam(6500, 2.5e-6).rms_emittance == 3.6082179183888383e-10
+
+
+class TestBeamCalculations:
+ def test_lhc_revolution_frequency(self):
+ lhc_beam = Beam(6500, 2.5e-6)
+ assert lhc_beam.revolution_frequency() == 11245.499628523643
+
+ @pytest.mark.parametrize(
+ "alpha_p, result",
+ [(0, 2.083077890845299e-08), (1e-5, -9.979169221091548e-06), (-500, 500.0000000208308)],
+ )
+ def test_eta(self, alpha_p, result):
+ assert Beam(6500, 2.5e-6).eta(alpha_p) == result
+
+ @pytest.mark.parametrize(
+ "alpha_p, result", [(1e-5, 316.2277660168379), (500, 0.044721359549995794)],
+ )
+ def test_gamma_transition(self, alpha_p, result):
+ assert Beam(6500, 2.5e-6).gamma_transition(alpha_p) == result
+
+ def test_gamma_transition_raises(self):
+ with pytest.raises(ZeroDivisionError):
+ Beam(6500, 2.5e-6).gamma_transition(0)
+
+
+class TestRipken:
+ def test_beam_size(self, _fake_coordinates):
+ # Test uses the dispatcher's 'py_func' object for coverage integrity, packaged implementation
+ # is still JIT compiled, and a JIT failure will fail the test.
+ assert np.allclose(ripken._beam_size.py_func(_fake_coordinates), _fake_coordinates.std())
+ assert np.allclose(
+ ripken._beam_size.py_func(_fake_coordinates, method="rms"),
+ np.sqrt(np.mean(np.square(_fake_coordinates))),
+ )
+
+ def test_beam_size_raises(self, _fake_coordinates):
+ with pytest.raises(NotImplementedError):
+ _ = ripken._beam_size.py_func(_fake_coordinates, method="not_real")
+
+ @pytest.mark.parametrize("beta11", [0.3312])
+ @pytest.mark.parametrize("beta21", [1])
+ @pytest.mark.parametrize("emit_x", [5e-6, 2.75e-6, 3.5e-6])
+ @pytest.mark.parametrize("emit_y", [5e-6, 2.75e-6, 3.5e-6])
+ def test_lebedev_size_floats(self, beta11, beta21, emit_x, emit_y):
+ assert ripken.lebedev_beam_size(
+ beta1_=beta11, beta2_=beta21, geom_emit_x=emit_x, geom_emit_y=emit_y
+ ) == np.sqrt(emit_x * beta11 + emit_y * beta21)
+
+
+# TODO: fix this one
+# def test_levedev_size_arrays(self):
+# beta_11_array = np.load(INPUT_PATHS["beta11"])
+# beta_21_array = np.load(INPUT_PATHS["beta21"])
+# lebedev_sizes = np.load(INPUT_PATHS["lebedev"])
+# calculated = ripken.lebedev_beam_size(beta_11_array, beta_21_array, 3.75e-6, 3.75e-6)
+# assert np.allclose(calculated, lebedev_sizes)
+
+
+class TestTwiss:
+ def test_courant_snyder_transform(self):
+ # Test uses the dispatcher's 'py_func' object for coverage integrity, packaged implementation
+ # is still JIT compiled, and a JIT failure will fail the test.
+ alpha_beta = np.load(INPUT_PATHS["alpha_beta"])
+ u_vector = np.load(INPUT_PATHS["u_vector"])
+ u_bar_result = np.load(INPUT_PATHS["u_bar"])
+ u_transform = twiss.courant_snyder_transform.py_func(u_vector, alpha_beta[0], alpha_beta[1])
+ np.testing.assert_array_almost_equal(u_transform, u_bar_result)
+
+
+@pytest.fixture()
+def _fake_coordinates() -> np.ndarray:
+ return np.random.random(size=10_000) / 1e4
diff --git a/tests/test_plotting.py b/tests/test_plotting.py
deleted file mode 100644
index ae0c1fc0..00000000
--- a/tests/test_plotting.py
+++ /dev/null
@@ -1,27 +0,0 @@
-import matplotlib
-import matplotlib.pyplot as plt
-import pytest
-
-from pyhdtoolkit.plotting.helpers import AnnotationsPlotter
-
-# Forcing non-interactive Agg backend so rendering is done similarly across platforms during tests
-matplotlib.use("Agg")
-
-
-class TestAnnotationsPlotter:
- @pytest.mark.mpl_image_compare(tolerance=20, style="seaborn-pastel", savefig_kwargs={"dpi": 200})
- def test_set_arrow_label(self):
- figure = plt.figure(figsize=(12, 7))
- ax = figure.add_subplot()
- ax.scatter(0, 0)
- ax.set_xlim((-1, 2))
- ax.set_ylim((-1, 2))
- AnnotationsPlotter.set_arrow_label(
- axis=ax,
- label="Test label on a test point",
- color="b",
- arrow_position=(0, 0),
- label_position=(1, 1),
- arrow_arc_rad=-0.2,
- )
- return figure
diff --git a/tests/test_twiss.py b/tests/test_twiss.py
deleted file mode 100644
index bf5095cd..00000000
--- a/tests/test_twiss.py
+++ /dev/null
@@ -1,23 +0,0 @@
-import pathlib
-
-import numpy as np
-
-from pyhdtoolkit.optics import twiss
-
-CURRENT_DIR = pathlib.Path(__file__).parent
-INPUTS_DIR = CURRENT_DIR / "inputs"
-INPUT_PATHS = {
- "alpha_beta": INPUTS_DIR / "alpha_beta.npy",
- "u_vector": INPUTS_DIR / "u_vector.npy",
- "u_bar": INPUTS_DIR / "u_bar.npy",
-}
-
-
-def test_courant_snyder_transform():
- # Test uses the dispatcher's 'py_func' object for coverage integrity, packaged implementation
- # is still JIT compiled, and a JIT failure will fail the test.
- alpha_beta = np.load(INPUT_PATHS["alpha_beta"])
- u_vector = np.load(INPUT_PATHS["u_vector"])
- u_bar_result = np.load(INPUT_PATHS["u_bar"])
- u_transform = twiss.courant_snyder_transform.py_func(u_vector, alpha_beta[0], alpha_beta[1])
- np.testing.assert_array_almost_equal(u_transform, u_bar_result)
diff --git a/tests/test_utils.py b/tests/test_utils.py
index ebdc3ca7..e25bff08 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -7,6 +7,8 @@
import pytest
+from loguru import logger
+
from pyhdtoolkit.utils import defaults # here for coverage
from pyhdtoolkit.utils.cmdline import CommandLine
from pyhdtoolkit.utils.executors import MultiProcessor, MultiThreader
@@ -26,6 +28,19 @@ def _to_str(integer: int) -> str:
return str(integer)
+class TestDefaults:
+ def test_logger_config(self, capsys):
+ defaults.config_logger()
+ message = "This should be in stdout now"
+ logger.info(message)
+ captured = capsys.readouterr()
+ assert message in captured.out
+
+ # This is to get it back as it is by defaults for other tests
+ logger.remove()
+ logger.add(sys.stderr)
+
+
@pytest.mark.skipif(
sys.platform.startswith("win"), reason="Windows is a shitshow for this.",
)