From 93777bbf60b664ee5e2985e39a15cdfdedf1d4dc Mon Sep 17 00:00:00 2001 From: Liam Pattinson Date: Tue, 4 Jul 2023 17:43:46 +0100 Subject: [PATCH 1/4] Add COCOS features, initial Python package bits --- .flake8 | 4 + .github/workflows/publish.yml | 26 ++++ .github/workflows/test.yml | 31 +++++ .gitignore | 3 + LICENSE | 2 +- README.md | 35 ++++++ pyloidal/__init__.py | 5 + pyloidal/cocos.py | 221 ++++++++++++++++++++++++++++++++++ pyproject.toml | 47 ++++++++ tests/test_cocos.py | 67 +++++++++++ 10 files changed, 440 insertions(+), 1 deletion(-) create mode 100644 .flake8 create mode 100644 .github/workflows/publish.yml create mode 100644 .github/workflows/test.yml create mode 100644 pyloidal/__init__.py create mode 100644 pyloidal/cocos.py create mode 100644 pyproject.toml create mode 100644 tests/test_cocos.py diff --git a/.flake8 b/.flake8 new file mode 100644 index 0000000..c5c6294 --- /dev/null +++ b/.flake8 @@ -0,0 +1,4 @@ +[flake8] +max-line-length = 88 +extend-ignore = E203 + diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 0000000..356e9ec --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,26 @@ +name: Publish to PyPI + +on: + release: + types: [published] + +jobs: + deploy: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: '3.x' + - name: Install + run: | + python -m pip install --upgrade pip + python -m pip install build + - name: Build package + run: python -m build --sdist --wheel --outdir dist/ + - name: Publish package + uses: pypa/gh-action-pypi-publish@release/v1 + with: + user: __token__ + password: ${{ secrets.PYPI_TOKEN }} diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..4f57767 --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,31 @@ +name: Tests + +on: + push: + paths: + - '**.py' + pull_request: + paths: + - '**.py' + +jobs: + run_tests: + name: Run tests + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.8", "3.9", "3.10"] + steps: + - uses: actions/checkout@v3 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + - name: Install + run: | + python -m pip install --upgrade pip + python -m pip install .[test] + - name: Test + run: | + pytest -vv diff --git a/.gitignore b/.gitignore index 68bc17f..21e5cd2 100644 --- a/.gitignore +++ b/.gitignore @@ -158,3 +158,6 @@ cython_debug/ # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ + +# Setuptools scm _version.py file +_version.py diff --git a/LICENSE b/LICENSE index be901a1..374ee79 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ MIT License -Copyright (c) 2023 PlasmaFAIR +Copyright (c) 2023 PlasmaFAIR, 2017 Orso Meneghini Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/README.md b/README.md index 302140e..e700989 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,37 @@ # pyloidal + Python utilities for tokamak science. + +## Install + +Install from PyPI: + +```bash +$ python3 -m pip install --upgrade pip +$ python3 -m pip install pyloidal +``` + +Install from GitHub repo: + +```bash +$ git clone https://github.com/PlasmaFAIR/pyloidal +$ cd pyloidal +$ python3 -m pip install --upgrade pip +$ python3 -m pip install . +``` + +## Tests + +First clone the repo, then: + +```bash +$ python3 -m pip install .[test] +$ pytest +``` + +## License + +Copyrighted under the MIT License. Uses elements from [OMAS][OMAS], which is also +copyrighted under MIT, 2017 Orso Meneghini. + +[OMAS]: https://github.com/gafusion/omas diff --git a/pyloidal/__init__.py b/pyloidal/__init__.py new file mode 100644 index 0000000..6044291 --- /dev/null +++ b/pyloidal/__init__.py @@ -0,0 +1,5 @@ +from contextlib import suppress +from importlib.metadata import PackageNotFoundError, version + +with suppress(PackageNotFoundError): + __version__ = version("pyloidal") diff --git a/pyloidal/cocos.py b/pyloidal/cocos.py new file mode 100644 index 0000000..fea3dd0 --- /dev/null +++ b/pyloidal/cocos.py @@ -0,0 +1,221 @@ +r""" +Determine the COCOS (COordinate COnventionS) convention used to describe a tokamak +equilibrium and convert between COCOS systems. Based on the work of Sauter, O. and +Medvedev, S.Y., 2013. Tokamak coordinate conventions: COCOS. *Computer Physics +Communications*, 184(2), pp.293-302. + +Throughout, we denote the coordinate systems in a tokamak with the following terms: + +- :math:`R`, the major radius +- :math:`r`, the minor radius +- :math:`Z`, the vertical coordinate +- :math:`\phi`, the toroidal angle +- :math:`\theta`, the poloidal angle + +These functions were adapted from OMAS (Copyright MIT License, 2017, Orso Meneghini). +""" + +import itertools +from typing import Dict, Optional, Tuple, Union + +import numpy as np +from numpy.typing import ArrayLike + + +def sigma_to_cocos( + sigma_bp: int, + sigma_rpz: int, + sigma_rtp: int, + psi_by_2pi: bool = True, +) -> int: + r""" + We can (partially) determine the COCOS by checking the :math:`\sigma` quantities, + defined below. We additionally need to know whether :math:`\psi` is defined with a + factor of :math:`1/(2\pi)` to determine whether the COCOS is in the range 1 to 8 or + 11 to 18. + + Parameters + ---------- + sigma_rtp + :math:`\sigma_{B_p}`: Given by :math:`\text{sign}(\vec{B}_p`\cdot\nabla\theta)`, + so is +1 when the poloidal magnetic field is in the same direction as increasing + :math:`\theta` and -1 when they are opposed. + sigma_rpz + :math:`\sigma_{R \phi Z`}: +1 when :math:`(R, \phi`, Z)` form a right-handed + coordinate system, and -1 when :math:`(R, Z, \phi)` form a right-handed + coordinate system. + sigma_bp + :math:`\sigma_{r\theta\phi}`: +1 when :math:`(r, \theta, \phi)` form a + right-handed coordinate system, and -1 when :math:`(r, \theta, \phi)` form a + right-handed coordainte system. + psi_by_2pi + If true, :math:`\psi` is defined with a factor of :math:`1/(2\pi)`. + + Returns + ------- + int + COCOS convention in use. + """ + if sigma_bp != 1 and sigma_bp != -1: + raise ValueError(f"sigma_bp should be either 1 or -1, found {sigma_bp}") + if sigma_rpz != 1 and sigma_rpz != -1: + raise ValueError(f"sigma_rpz should be either 1 or -1, found {sigma_rpz}") + if sigma_rtp != 1 and sigma_rtp != -1: + raise ValueError(f"sigma_rtp should be either 1 or -1, found {sigma_rtp}") + + sigma_to_cocos_dict = { + (+1, +1, +1): 1, # +Bp, +rpz, +rtp + (+1, -1, +1): 2, # +Bp, -rpz, +rtp + (-1, +1, -1): 3, # -Bp, +rpz, -rtp + (-1, -1, -1): 4, # -Bp, -rpz, -rtp + (+1, +1, -1): 5, # +Bp, +rpz, -rtp + (+1, -1, -1): 6, # +Bp, -rpz, -rtp + (-1, +1, +1): 7, # -Bp, +rpz, +rtp + (-1, -1, +1): 8, # -Bp, -rpz, +rtp + } + result = sigma_to_cocos_dict[(sigma_bp, sigma_rpz, sigma_rtp)] + return result if psi_by_2pi else result + 10 + + +def identify_cocos( + b_toroidal: float, + plasma_current: float, + safety_factor: ArrayLike, + poloidal_flux: ArrayLike, + clockwise_phi: Optional[bool] = None, + minor_radii: Optional[ArrayLike] = None, +) -> Tuple[int, ...]: + r""" + Determine which COCOS coordinate system is in use. Returns all possible conventions. + + Parameters + ---------- + b_toroidal + Toroidal magnetic field, with sign. Should be in units of Tesla. + plasma_current + Plasma current, with sign. Should be in units of Amperes. + safety_factor + Safety factor profile, with sign, as a function of ``poloidal_flux``. Usually + denoted :math:`q`. Should be defined with the first element on the magnetic + axis, and subsequent elements on successively larger flux surfaces. + poloidal_flux + The profile of the poloidal flux function, with sign. Usually denoted + :math:`\psi`. Should be defined with the first element on the magnetic + axis, and subsequent elements on successively larger flux surfaces. + clockwise_phi + When viewing tokamak coordinates from above, does the toroidal angle + :math:`\phi` increase in the clockwise direction? This is required to identify + odd vs even COCOS. This cannot be determined from the output of a code alone. + An easy way to determine this is to answer the question: is positive + :math:`B_\text{toroidal}` clockwise? + minor_radii + Th minor radius of each flux surface as function of ``poloidal_flux``. This is + required to identify whether :math:`\psi` contains a factor of :math:`2\pi`. + + Returns + ------- + Tuple[int, ...] + All possible conventions are returned. If both optional arguments are provided, + the returned tuple should have length 1. + + """ + + if clockwise_phi is None: + return tuple( + itertools.chain.from_iterable( + identify_cocos( + b_toroidal, + plasma_current, + safety_factor, + poloidal_flux, + x, + minor_radii, + ) + for x in (False, False) + ) + ) + + sign_plasma_current = np.sign(plasma_current) + sign_b_toroidal = np.sign(b_toroidal) + sign_q = np.sign(safety_factor[0]) + psi_increasing = np.sign(poloidal_flux[1] - poloidal_flux[0]) + + sigma_bp = psi_increasing * sign_plasma_current + sigma_rpz = -1 if clockwise_phi else 1 + sigma_rtp = sign_q * sign_plasma_current * sign_b_toroidal + + # identify 2*pi term in poloidal_flux definition based on safety_factor estimate + if minor_radii is None: + # Return both variants if not provided with minor radii + return tuple( + itertools.chain.from_iterable( + sigma_to_cocos(sigma_bp, sigma_rpz, sigma_rtp, phi_by_2pi=x) + for x in (True, False) + ) + ) + + index = np.argmin(np.abs(safety_factor)) + if index == 0: + index += 1 + safety_factor_estimate = np.abs( + (np.pi * b_toroidal * (minor_radii[index] - minor_radii[0]) ** 2) + / (poloidal_flux[index] - poloidal_flux[0]) + ) + safety_factor_actual = np.abs(safety_factor[index]) + psi_by_2pi = np.abs( + safety_factor_estimate / (2 * np.pi) - safety_factor_actual + ) < np.abs(safety_factor_estimate - safety_factor_actual) + + return (sigma_to_cocos(sigma_bp, sigma_rpz, sigma_rtp, psi_by_2pi=psi_by_2pi),) + + +def cocos_coefficients(cocos: int) -> Dict[str, int]: + r"""Returns dictionary with COCOS coefficients given a COCOS index""" + # TODO Maybe have a pandas Dataframe with all the info written explicitly? + coeffs = { + "exp_bp": int(cocos >= 10), + "sigma_bp": 1 if cocos in (1, 2, 5, 6, 11, 12, 15, 16) else -1, + "sigma_rpz": 1 if cocos % 2 else -1, + "sigma_rtp": 1 if cocos in (1, 2, 7, 8, 11, 12, 17, 18) else -1, + } + coeffs["phi_clockwise"] = coeffs["sigma_rpz"] == -1 + coeffs["theta_clockwise"] = cocos in (1, 4, 6, 7, 11, 14, 16, 17) + coeffs["psi_increasing"] = bool(coeffs["exp_bp"]) + coeffs["sign_q"] = coeffs["sigma_rtp"] + coeffs["sign_pprime"] = -coeffs["sigma_bp"] + return coeffs + + +def cocos_transform(cocos_in: int, cocos_out: int) -> Dict[str, Union[int, float]]: + r""" + Returns a dictionary with coefficients for how various quantities should by + multiplied in order to go from ``cocos_in`` to ``cocos_out``. + """ + + coeffs_in = cocos_coefficients(cocos_in) + coeffs_out = cocos_coefficients(cocos_out) + + sigma_ip_eff = coeffs_in["sigma_rpz"] * coeffs_out["sigma_rpz"] + sigma_b0_eff = coeffs_in["sigma_rpz"] * coeffs_out["sigma_rpz"] + sigma_bp_eff = coeffs_in["sigma_bp"] * coeffs_out["sigma_bp"] + exp_bp_eff = coeffs_out["exp_bp"] - coeffs_in["exp_bp"] + sigma_rtp_eff = coeffs_in["sigma_rtp"] * coeffs_out["sigma_rtp"] + + # Transform + transforms = {} + transforms["1/psi"] = sigma_ip_eff * sigma_bp_eff / (2 * np.pi) ** exp_bp_eff + transforms["d/dpsi"] = transforms["1/psi"] + transforms["ffprime"] = transforms["1/psi"] + transforms["pprime"] = transforms["1/psi"] + + transforms["toroidal"] = sigma_b0_eff + transforms["b_toroidal"] = transforms["toroidal"] + transforms["plasma_current"] = transforms["toroidal"] + transforms["f"] = transforms["toroidal"] + + transforms["poloidal"] = sigma_b0_eff * sigma_rtp_eff + transforms["b_poloidal"] = transforms["poloidal"] + + transforms["psi"] = sigma_ip_eff * sigma_bp_eff * (2 * np.pi) ** exp_bp_eff + transforms["q"] = sigma_ip_eff * sigma_b0_eff * sigma_rtp_eff + return transforms diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..a95c12e --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,47 @@ +[project] +name = "pyloidal" +description = "Python utilies for tokamak science." +readme = "README.md" +authors = [ + {name = "Liam Pattinson", email = "liam.pattinson@york.ac.uk"} +] +license = {file = "LICENSE"} +dynamic = ["version"] +keywords = ["tokamak", "cocos"] +classifiers = [ + "Programming Language :: Python", + "Development Status :: 3 - Alpha", + "Natural Language :: English", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent", + "Programming Language :: Python", + "Topic :: Software Development :: Libraries :: Python Modules", + "Topic :: Scientific/Engineering :: Physics", +] + +requires-python = ">=3.8" +dependencies = ["numpy ~= 1.24"] + +[project.optional-dependencies] +test = ["pytest"] +lint = ["black", "flake8", "isort", "ruff"] + +[project.urls] +Source = "https://github.com/PlasmaFAIR/pyloidal" +Tracker = "https://github.com/PlasmaFAIR/pyloidal/issues" + +[build-system] +requires = [ + "setuptools >= 65", + "setuptools_scm[toml] >= 6.2", +] +build-backend = "setuptools.build_meta" + +[tool.setuptools_scm] +write_to = "pyloidal/_version.py" +fallback_version = "0.1.0" + +[tool.isort] +profile = "black" + diff --git a/tests/test_cocos.py b/tests/test_cocos.py new file mode 100644 index 0000000..c18c8fd --- /dev/null +++ b/tests/test_cocos.py @@ -0,0 +1,67 @@ +import numpy as np +import pytest + +from pyloidal.cocos import identify_cocos + +# Create identify_cocos test cases for COCOS 1, 3, 5, 7 +# TODO include tests for negative/antiparallel b_toroidal and plasma_current +# TODO include tests for multiple returns + +identify_cocos_tests = {} +odds = { + 1: { + "b_toroidal": 2.5, + "plasma_current": 1e6, + "poloidal_flux": np.linspace(0, 2, 3), + "safety_factor": np.linspace(0.5, 1.5, 3), + "clockwise_phi": False, + "minor_radii": np.linspace(0, 2, 3), + }, + 3: { + "b_toroidal": 2.5, + "plasma_current": 1e6, + "poloidal_flux": np.linspace(2, 0, 3), + "safety_factor": np.linspace(-0.5, -1.5, 3), + "clockwise_phi": False, + "minor_radii": np.linspace(0, 2, 3), + }, + 5: { + "b_toroidal": 2.5, + "plasma_current": 1e6, + "poloidal_flux": np.linspace(0, 2, 3), + "safety_factor": np.linspace(-0.5, -1.5, 3), + "clockwise_phi": False, + "minor_radii": np.linspace(0, 2, 3), + }, + 7: { + "b_toroidal": 2.5, + "plasma_current": 1e6, + "poloidal_flux": np.linspace(2, 0, 3), + "safety_factor": np.linspace(0.5, 1.5, 3), + "clockwise_phi": False, + "minor_radii": np.linspace(0, 2, 3), + }, +} +identify_cocos_tests.update(odds) + +# Set clockwise_phi to True in these to get COCOS 2, 4, 6, 8 +evens = {} +for cocos, kwargs in odds.items(): + even_kwargs = kwargs.copy() + even_kwargs["clockwise_phi"] = True + evens[cocos + 1] = even_kwargs +identify_cocos_tests.update(evens) +# Multiply by factor of 2*pi to get COCOS 11 -> 18 +tens = {} +for cocos, kwargs in identify_cocos_tests.items(): + tens_kwargs = kwargs.copy() + # Note: can't use *= here, as some references are shared between tests + tens_kwargs["poloidal_flux"] = kwargs["poloidal_flux"] * (2 * np.pi) + tens_kwargs["safety_factor"] = kwargs["safety_factor"] * (2 * np.pi) + tens[cocos + 10] = tens_kwargs +identify_cocos_tests.update(tens) + + +@pytest.mark.parametrize("expected_cocos,kwargs", [*identify_cocos_tests.items()]) +def test_identify_cocos(expected_cocos, kwargs): + assert identify_cocos(**kwargs) == (expected_cocos,) From 4dd558cb281f3d1b0350a1828bc60b9dcdc93160 Mon Sep 17 00:00:00 2001 From: Liam Pattinson Date: Tue, 4 Jul 2023 17:56:41 +0100 Subject: [PATCH 2/4] Add cocos_transform test --- tests/test_cocos.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/test_cocos.py b/tests/test_cocos.py index c18c8fd..8752cdd 100644 --- a/tests/test_cocos.py +++ b/tests/test_cocos.py @@ -1,7 +1,7 @@ import numpy as np import pytest -from pyloidal.cocos import identify_cocos +from pyloidal.cocos import identify_cocos, cocos_transform # Create identify_cocos test cases for COCOS 1, 3, 5, 7 # TODO include tests for negative/antiparallel b_toroidal and plasma_current @@ -65,3 +65,12 @@ @pytest.mark.parametrize("expected_cocos,kwargs", [*identify_cocos_tests.items()]) def test_identify_cocos(expected_cocos, kwargs): assert identify_cocos(**kwargs) == (expected_cocos,) + +def test_cocos_transform(): + # TODO should be parametrized + assert cocos_transform(1, 3)['poloidal'] == -1 + for cocos in range(1, 9): + assert cocos_transform(cocos, cocos + 10)['1/psi'] != 1 + for cocos_add in (cocos + x * 10 for x in range(2)): + for key in ('b_toroidal', 'toroidal', 'poloidal', 'q'): + assert cocos_transform(cocos_add, cocos_add)[key] == 1 From 7cd16d4a57855f4ad3a34dc14509cf34a186c1f8 Mon Sep 17 00:00:00 2001 From: LiamPattinson Date: Tue, 4 Jul 2023 18:08:12 +0100 Subject: [PATCH 3/4] Fix bug in identify_cocos Co-authored-by: Peter Hill --- pyloidal/cocos.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyloidal/cocos.py b/pyloidal/cocos.py index fea3dd0..1871b0c 100644 --- a/pyloidal/cocos.py +++ b/pyloidal/cocos.py @@ -131,7 +131,7 @@ def identify_cocos( x, minor_radii, ) - for x in (False, False) + for x in (True, False) ) ) From e1f28023f954919e943c2699eb05d5277d9f3145 Mon Sep 17 00:00:00 2001 From: Liam Pattinson Date: Tue, 4 Jul 2023 18:11:19 +0100 Subject: [PATCH 4/4] Add Python 3.11 to test workflow --- .github/workflows/test.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 4f57767..866c69a 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -15,7 +15,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.8", "3.9", "3.10"] + python-version: ["3.8", "3.9", "3.10", "3.11"] steps: - uses: actions/checkout@v3 - name: Set up Python ${{ matrix.python-version }}