diff --git a/astar_utils/__init__.py b/astar_utils/__init__.py index f0a9d24..a6c364d 100644 --- a/astar_utils/__init__.py +++ b/astar_utils/__init__.py @@ -5,3 +5,4 @@ from .unique_list import UniqueList from .badges import Badge, BadgeReport from .loggers import get_logger, get_astar_logger +from .spectral_types import SpectralType diff --git a/astar_utils/spectral_types.py b/astar_utils/spectral_types.py new file mode 100644 index 0000000..95feb43 --- /dev/null +++ b/astar_utils/spectral_types.py @@ -0,0 +1,161 @@ +# -*- coding: utf-8 -*- +"""Contains SpectralType class.""" + +import re +from typing import ClassVar +from dataclasses import dataclass, field, InitVar + + +@dataclass(frozen=True, slots=True) +class SpectralType: + r"""Parse and store stellar spectral types. + + This dataclass can understand a string constructor containing a valid + stellar spectral type including "OBAFGKM"-spectral class, 0-9 subclass + (including floats up to one decimal place) and luminosity class (roman + numerals I-V), which are converted to attributes. The initial spectral + type must match the regex "^[OBAFGKM](\d(\.\d)?)?(I{1,3}|IV|V)?$". Only the + main spectral class ("OBAFGKM") is mandatory, both numerical subclass and + luminosity class may be empty (None). Spectral subclasses (0.0-9.9) are + internally stored as floats, meaning an initial type of e.g. "A0V" is + considered identical to "A0.0V". Representations of an instance (``str`` + and ``repr``) always show integers whenever possible, regardless of the + initially passed value. + + Instances of this class are always "frozen", meaning they cannot be changed + after initial creation. It is not possible to manually assign the three + individual attribute, creation can only happen via the constructor string + described above. + + One of the main features of this class (and indeed motivation for its + creation in the first place) is the ability to compare, and thus order, + instances of the class based on their two-component spectral (sub)class. + In this context, a "later" spectral type is considered "greater" than an + "earlier" one, i.e. O < B < A < F < G < K < M, which is consistent with the + convention of numerical subtypes, where 0 < 9 holds true. This is, however, + in contrast to the physical meaning of this parameter, which is correlated + with stellar effective temperature in reverse order, meaning T(A) > T(G) + and T(F0) > T(F9), by convention. On the other hand, in many visualisations + such as the Herzsprung-Russel diagram, it is common practice to represent + temperature on the x-axis in descending order, meaning a sorted list of + instances of this class will already have the typical order of such + diagrams. + + In this context, the luminosity class (if any) is ignored for sorting and + comparison (<, >, <=, >=), as it represents a second physical dimension. + However, instances of this class may also be compared for equality (== and + !=), in which case all three attributes are considered. + + Attributes + ---------- + spectral_class : str + Main spectral class (OBAFGKM). + spectral_subclass : str or None + Numerical spectral subclass (0.0-9.9). + luminosity_class : str or None + Roman numeral luminosity class (I-V). + + Notes + ----- + The constructor string can be supplied in both upper or lower case or a + mixture thereof, meaning "A0V", "a0v", "A0v" and "a0V" are all valid + representations of the same spectral type. The internal attributes are + converted to uppercase upon creation. + + Examples + -------- + >>> from astar_utils import SpectralType + >>> spt = SpectralType("A0V") + >>> spt.spectral_class + 'A' + + >>> spt.spectral_subclass + 0.0 + + >>> spt.luminosity_class + 'V' + + >>> spts = [SpectralType(s) for s in + ... ["G2", "M4.0", "B3", "F8", "K6.5", + ... "A", "A0", "A9", "O8"]] + >>> sorted(spts) # doctest: +NORMALIZE_WHITESPACE + [SpectralType('O8'), + SpectralType('B3'), + SpectralType('A0'), + SpectralType('A'), + SpectralType('A9'), + SpectralType('F8'), + SpectralType('G2'), + SpectralType('K6.5'), + SpectralType('M4')] + + Note that a missing spectral subtype is considered as 5 (middle of the + main spectral class) in the context of sorting, as shown in the example + with the spectral type "A" ending up between "A0" and "A9". + """ + + spectral_class: str = field(init=False, default="") + spectral_subclass: float | None = field(init=False, default=None) + luminosity_class: str | None = field(init=False, default=None) + spectype: InitVar[str] + _cls_order: ClassVar = "OBAFGKM" # descending Teff + _regex: ClassVar = re.compile( + r"^(?P[OBAFGKM])(?P\d(\.\d)?)?" + "(?PI{1,3}|IV|V)?$", re.ASCII | re.IGNORECASE) + + def __post_init__(self, spectype) -> None: + """Validate input and populate fields.""" + if not (match := self._regex.fullmatch(spectype)): + raise ValueError(spectype) + + classes = match.groupdict() + # Circumvent frozen as per the docs... + object.__setattr__(self, "spectral_class", + str(classes["spec_cls"]).upper()) + + if classes["sub_cls"] is not None: + object.__setattr__(self, "spectral_subclass", + float(classes["sub_cls"])) + + if classes["lum_cls"] is not None: + object.__setattr__(self, "luminosity_class", + str(classes["lum_cls"]).upper()) + + @property + def _subcls_str(self) -> str: + if self.spectral_subclass is None: + return "" + if self.spectral_subclass.is_integer(): + return str(int(self.spectral_subclass)) + return str(self.spectral_subclass) + + def __repr__(self) -> str: + """Return repr(self).""" + return f"{self.__class__.__name__}('{self!s}')" + + def __str__(self) -> str: + """Return str(self).""" + spectype = (f"{self.spectral_class}{self._subcls_str}" + f"{self.luminosity_class or ''}") + return spectype + + @property + def _comp_tuple(self) -> tuple[int, float]: + # if None, assume middle of spectral class + if self.spectral_subclass is not None: + sub_cls = self.spectral_subclass + else: + sub_cls = 5 + return (self._cls_order.index(self.spectral_class), sub_cls) + + def __lt__(self, other) -> bool: + """Return self < other.""" + if not isinstance(other, self.__class__): + raise TypeError("Can only compare equal types.") + return self._comp_tuple < other._comp_tuple + + def __le__(self, other) -> bool: + """Return self < other.""" + if not isinstance(other, self.__class__): + raise TypeError("Can only compare equal types.") + return self._comp_tuple <= other._comp_tuple diff --git a/pyproject.toml b/pyproject.toml index ccab792..c0321c1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "astar-utils" -version = "0.3.1a0" +version = "0.3.1a1" description = "Contains commonly-used utilities for AstarVienna's projects." license = "GPL-3.0-or-later" authors = ["Fabian Haberhauer "] diff --git a/tests/test_spectral_types.py b/tests/test_spectral_types.py new file mode 100644 index 0000000..02e534e --- /dev/null +++ b/tests/test_spectral_types.py @@ -0,0 +1,145 @@ +# -*- coding: utf-8 -*- +"""Unit tests for spectral_types.py.""" + +import pytest +import operator + +from astar_utils import SpectralType + + +class TestParsesTypes: + @pytest.mark.parametrize("spec_cls", [*"OBAFGKM"]) + def test_parses_valid_spec_cls(self, spec_cls): + spt = SpectralType(spec_cls) + assert spt.spectral_class == spec_cls + assert spt.spectral_subclass is None + assert spt.luminosity_class is None + + def test_fails_on_invalid_spec_cls(self): + with pytest.raises(ValueError): + SpectralType("X") + + @pytest.mark.parametrize(("spec_cls", "sub_cls"), + [(*"A0",), (*"G2",), ("F", "3.5"), ("K", "7.5"), + (*"M2",), (*"B9",), ("G", "5.0"), ("F", "0.5")]) + def test_parses_valid_spec_and_sub_cls(self, spec_cls, sub_cls): + spt = SpectralType(spec_cls + sub_cls) + assert spt.spectral_class == spec_cls + assert spt.spectral_subclass == float(sub_cls) + assert spt.luminosity_class is None + + @pytest.mark.parametrize("spec_sub_cls", ["X2", "Y3.5", "G99", "K1.2222"]) + def test_fails_on_invalid_spec_and_sub_cls(self, spec_sub_cls): + with pytest.raises(ValueError): + SpectralType(spec_sub_cls) + + @pytest.mark.parametrize(("spec_cls", "lum_cls"), + [(*"AV",), (*"GI",), ("F", "III"), ("K", "IV")]) + def test_parses_valid_spec_and_lum_cls(self, spec_cls, lum_cls): + spt = SpectralType(spec_cls + lum_cls) + assert spt.spectral_class == spec_cls + assert spt.spectral_subclass is None + assert spt.luminosity_class == lum_cls + + @pytest.mark.parametrize("spec_lum_cls", ["II", "YVII", "GIIV", "KIVI"]) + def test_fails_on_invalid_spec_and_lum_cls(self, spec_lum_cls): + with pytest.raises(ValueError): + SpectralType(spec_lum_cls) + + @pytest.mark.parametrize(("spec_cls", "sub_cls", "lum_cls"), + [(*"A0V",), (*"G2I",), ("F", "3.5", "II"), + ("K", "7.5", "IV"), (*"M2V",), (*"B9I",), + ("G", "5.0", "III"), ("F", "0.5", "V")]) + def test_parses_valid_spec_sub_lum_cls(self, spec_cls, sub_cls, lum_cls): + spt = SpectralType(spec_cls + sub_cls + lum_cls) + assert spt.spectral_class == spec_cls + assert spt.spectral_subclass == float(sub_cls) + assert spt.luminosity_class == lum_cls + + @pytest.mark.parametrize("spec_sub_lum_cls", + ["Bogus", "G3IIV", "K4.5IVI", "B999IV"]) + def test_fails_on_invalid_spec_sub_lum_cls(self, spec_sub_lum_cls): + with pytest.raises(ValueError): + SpectralType(spec_sub_lum_cls) + + +class TestComparesTypes: + @pytest.mark.parametrize(("ssl_cls_a", "ssl_cls_b"), + [("A0V", "A0V"), ("G2", "G2.0"), + ("M3III", "M3.0III")]) + def test_compares_classes_as_equal(self, ssl_cls_a, ssl_cls_b): + spt_a = SpectralType(ssl_cls_a) + spt_b = SpectralType(ssl_cls_b) + assert spt_a == spt_b + + @pytest.mark.parametrize(("ssl_cls_a", "ssl_cls_b"), + [("A0V", "A1V"), ("G2", "G2.5"), + ("M3III", "M8.0III"), ("B3V", "BV")]) + def test_compares_order_within_spec_cls(self, ssl_cls_a, ssl_cls_b): + spt_a = SpectralType(ssl_cls_a) + spt_b = SpectralType(ssl_cls_b) + assert spt_a < spt_b + assert spt_b > spt_a + + @pytest.mark.parametrize(("ssl_cls_a", "ssl_cls_b"), + [("A0V", "G1V"), ("G", "M"), ("OII", "GIV"), + ("K3III", "M8.0III"), ("BV", "KI")]) + def test_compares_order_across_spec_cls(self, ssl_cls_a, ssl_cls_b): + spt_a = SpectralType(ssl_cls_a) + spt_b = SpectralType(ssl_cls_b) + assert spt_a < spt_b + assert spt_b > spt_a + + @pytest.mark.parametrize(("ssl_cls_a", "ssl_cls_b"), + [("A0V", "G1V"), ("G", "M"), ("OII", "GIV"), + ("K3III", "M8.0III"), ("BV", "KI"), + ("A0V", "A1V"), ("G2", "G2.5"), + ("M3III", "M8.0III"), ("B3V", "BV")]) + def test_compares_order_with_equal(self, ssl_cls_a, ssl_cls_b): + spt_a = SpectralType(ssl_cls_a) + spt_b = SpectralType(ssl_cls_b) + assert spt_a <= spt_b + assert spt_b >= spt_a + + @pytest.mark.parametrize("operation", [operator.gt, operator.lt, + operator.ge, operator.le]) + def test_throws_on_invalid_compare(self, operation): + with pytest.raises(TypeError): + operation(SpectralType("A0V"), 42) + + +class TestRepresentations: + @pytest.mark.parametrize(("ssl_cls", "exptcted"), + [("A0V", "SpectralType('A0V')"), + ("G2", "SpectralType('G2')"), + ("K9.0", "SpectralType('K9')"), + ("B2.5", "SpectralType('B2.5')"), + ("M3.0III", "SpectralType('M3III')"), + ("KII", "SpectralType('KII')"),]) + def test_repr(self, ssl_cls, exptcted): + spt = SpectralType(ssl_cls) + assert repr(spt) == exptcted + + @pytest.mark.parametrize(("ssl_cls", "exptcted"), + [("A0V", "A0V"), + ("G2", "G2"), + ("K9.0", "K9"), + ("B2.5", "B2.5"), + ("M3.0III", "M3III"), + ("KII", "KII"),]) + def test_str(self, ssl_cls, exptcted): + spt = SpectralType(ssl_cls) + assert str(spt) == exptcted + + +class TestUpperLowerCase: + def test_uppers_lower_case(self): + spt_a = SpectralType("A0V") + spt_b = SpectralType("a0v") + assert spt_a == spt_b + + @pytest.mark.parametrize("mixcase", ["m2.5III", "M2.5iii"]) + def test_uppers_lower_case_mixed(self, mixcase): + spt_a = SpectralType("M2.5III") + spt_b = SpectralType(mixcase) + assert spt_a == spt_b