diff --git a/matplotlib_scalebar/dimension.py b/matplotlib_scalebar/dimension.py index c93fc08..b213a54 100644 --- a/matplotlib_scalebar/dimension.py +++ b/matplotlib_scalebar/dimension.py @@ -81,7 +81,6 @@ def calculate_preferred(self, value, units): if index: newunits, factor = units_factor[index - 1] return base_value / factor, newunits - else: return value, units @@ -166,3 +165,26 @@ def __init__(self): def create_label(self, value, latexrepr): # Overriden to remove space between value and units. return "{}{}".format(value, latexrepr) + + +class AtomicLengthDimension(SILengthDimension): + """Dimension for atomic-scale lengths using Angstrom as the base unit.""" + def __init__(self): + super().__init__() + latexrepr = "$\\mathrm{\\AA}$" + factor = 1e-10 + + self.add_units("angstrom", factor, latexrepr) # Full name + self.add_units("A", factor, latexrepr) + self.add_units("Å", factor, latexrepr) + +class AtomicLengthReciprocalDimension(SILengthReciprocalDimension): + """Dimension for reciprocal atomic-scale lengths using inverse Angstrom as the base unit.""" + def __init__(self): + super().__init__() + latexrepr = "$\\mathrm{\\AA}^{-1}$" + factor = 1e10 + + self.add_units("1/angstrom", factor, latexrepr) # Full name + self.add_units("1/A", factor, latexrepr) + self.add_units("1/Å", factor, latexrepr) \ No newline at end of file diff --git a/matplotlib_scalebar/scalebar.py b/matplotlib_scalebar/scalebar.py index efd5a5d..1409380 100644 --- a/matplotlib_scalebar/scalebar.py +++ b/matplotlib_scalebar/scalebar.py @@ -35,6 +35,8 @@ "IMPERIAL_LENGTH", "ASTRO_LENGTH", "PIXEL_LENGTH", + "ATOMIC_LENGTH", + "ATOMIC_LENGTH_RECIPROCAL", ] # Standard library modules. @@ -71,6 +73,8 @@ AstronomicalLengthDimension, PixelLengthDimension, AngleDimension, + AtomicLengthDimension, + AtomicLengthReciprocalDimension, ) # Globals and constants variables. @@ -132,6 +136,8 @@ def _validate_legend_loc(loc): ASTRO_LENGTH = "astro-length" PIXEL_LENGTH = "pixel-length" ANGLE = "angle" +ATOMIC_LENGTH = "atomic-length" +ATOMIC_LENGTH_RECIPROCAL = "atomic-length-reciprocal" _DIMENSION_LOOKUP = { SI_LENGTH: SILengthDimension, @@ -140,6 +146,8 @@ def _validate_legend_loc(loc): ASTRO_LENGTH: AstronomicalLengthDimension, PIXEL_LENGTH: PixelLengthDimension, ANGLE: AngleDimension, + ATOMIC_LENGTH: AtomicLengthDimension, + ATOMIC_LENGTH_RECIPROCAL: AtomicLengthReciprocalDimension, } @@ -229,6 +237,8 @@ def __init__( * ``:const:`astro-length```: scale bar showing pc, kpc ly, AU, etc. * ``:const:`pixel-length```: scale bar showing px, kpx, Mpx, etc. * ``:const:`angle```: scale bar showing \u00b0, \u2032 or \u2032\u2032. + * ``:const:`atomic-length```: scale bar showing \AA, nm, \u00B5m, etc. + * ``:const:`atomic-length-reciprocal```: scale bar showing 1/\AA, 1/nm, 1/\u00B5m, etc. * a :class:`matplotlib_scalebar.dimension._Dimension` object :type dimension: :class:`str` or :class:`matplotlib_scalebar.dimension._Dimension` @@ -601,8 +611,33 @@ def get_units(self): return self._units def set_units(self, units): + """ + Set the units of the scale bar. + """ + old_units = self._units if hasattr(self, '_units') else None + old_dx = self._dx if hasattr(self, '_dx') else None + + # Auto-detect Angstrom units and switch to atomic dimension + if units in ["Å", "A", "angstrom"]: + self._dimension = _DIMENSION_LOOKUP[ATOMIC_LENGTH]() + elif units in ["1/Å", "1/A", "1/angstrom"]: + self._dimension = _DIMENSION_LOOKUP[ATOMIC_LENGTH_RECIPROCAL]() + elif hasattr(self, '_dimension') and isinstance(self._dimension, (AtomicLengthDimension, AtomicLengthReciprocalDimension)) and units not in ["Å", "A", "angstrom", "1/Å", "1/A", "1/angstrom"]: + # Switch back to SI if changing from Angstrom to non-Angstrom units + if '/' in units: + self._dimension = _DIMENSION_LOOKUP[SI_LENGTH_RECIPROCAL]() + else: + self._dimension = _DIMENSION_LOOKUP[SI_LENGTH]() + if not self.dimension.is_valid_units(units): raise ValueError(f"Invalid unit ({units}) with dimension") + + # Convert dx to the new units if we're changing units + if old_units is not None and old_dx is not None: + # All conversions are now handled by the dimension's convert method + # since factors are relative to meters + self._dx = self.dimension.convert(old_dx, old_units, units) + self._units = units units = property(get_units, set_units) diff --git a/tests/test_dimension.py b/tests/test_dimension.py index b5d9362..69b26b0 100644 --- a/tests/test_dimension.py +++ b/tests/test_dimension.py @@ -11,6 +11,8 @@ SILengthReciprocalDimension, ImperialLengthDimension, PixelLengthDimension, + AtomicLengthDimension, + AtomicLengthReciprocalDimension, _LATEX_MU, ) @@ -37,6 +39,11 @@ (PixelLengthDimension(), 200, "px", 200.0, "px"), (PixelLengthDimension(), 0.02, "px", 0.02, "px"), (PixelLengthDimension(), 0.001, "px", 0.001, "px"), + # Test Angstrom preferred units + (AtomicLengthDimension(), 0.1, "nm", 1, "Å"), + (AtomicLengthDimension(), 10, "Å", 1, "nm"), + (AtomicLengthReciprocalDimension(), 100, "1/nm", 10, "1/Å"), + (AtomicLengthReciprocalDimension(), 0.1, "1/Å", 1, "1/nm"), ], ) def test_calculate_preferred(dim, value, units, expected_value, expected_units): @@ -65,6 +72,13 @@ def test_to_latex(dim, units, expected): (SILengthDimension(), 2, "um", "cm", 2e-4), (PixelLengthDimension(), 2, "kpx", "px", 2000), (PixelLengthDimension(), 2, "px", "kpx", 2e-3), + # Test Angstrom conversions + (AtomicLengthDimension(), 1, "Å", "nm", 0.1), + (AtomicLengthDimension(), 1, "nm", "Å", 10), + (AtomicLengthDimension(), 1, "A", "angstrom", 1), + (AtomicLengthReciprocalDimension(), 1, "1/Å", "1/nm", 10), + (AtomicLengthReciprocalDimension(), 1, "1/nm", "1/Å", 0.1), + (AtomicLengthReciprocalDimension(), 1, "1/A", "1/angstrom", 1), ], ) def test_convert(dim, value, units, newunits, expected_value):