Add pyrolite.geochem.ind.REY and pyrolite.geochem.pyrochem.REY

pyrolite/geochem/__init__.py

@@ -13,11 +13,7 @@
 from . import parse
 from . import transform
 from . import norm
-from .ind import (
-    __common_elements__,
-    __common_oxides__,
-    REE,
-)
+from .ind import __common_elements__, __common_oxides__, REE, REY
 from .ions import set_default_ionic_charges
 
 set_default_ionic_charges()
@@ -84,6 +80,21 @@ def list_REE(self):
         """
         return [i for i in REE() if i in self._obj.columns]
 
+    @property
+    def list_REY(self):
+        """
+        Get the subset of columns which are Rare Earth Element names.
+
+        Returns
+        --------
+        :class:`list`
+
+        Notes
+        -------
+        The returned list will reorder REE based on atomic number.
+        """
+        return [i for i in REY() if i in self._obj.columns]
+
     @property
     def list_oxides(self):
         """
@@ -134,6 +145,21 @@ def REE(self):
     def REE(self, df):
         self._obj.loc[:, self.list_REE] = df
 
+    @property
+    def REY(self):
+        """
+        Get a Rare Earth Element + Yttrium subset of a DataFrame.
+
+        Returns
+        --------
+        :class:`pandas.Dataframe`
+        """
+        return self._obj[self.list_REY]
+
+    @REE.setter
+    def REY(self, df):
+        self._obj.loc[:, self.list_REY] = df
+
     @property
     def oxides(self):
         """

pyrolite/geochem/ind.py

@@ -7,7 +7,6 @@
 * Incompatibility indexes for spider plot ordering.
 """
 import re
-import functools
 import numpy as np
 import pandas as pd
 import periodictable as pt
@@ -41,7 +40,6 @@ def _load_radii():
 ########################################################################################
 
 
-@functools.lru_cache(maxsize=4)  # cache outputs for speed
 def common_elements(cutoff=92, output="string", order=None, as_set=False):
     """
     Provides a list of elements up to a particular cutoff (by default including U).
@@ -90,7 +88,88 @@ def common_elements(cutoff=92, output="string", order=None, as_set=False):
         return elements
 
 
-@functools.lru_cache(maxsize=2)  # cache outputs for speed
+def REE(output="string", dropPm=True):
+    """
+    Provides a list of Rare Earth Elements.
+
+    Parameters
+    -----------
+    output : :class:`str`
+        Whether to return output list as formulae ('formula') or strings (anthing else).
+
+    Returns
+    -------
+    :class:`list` | :class:`set`
+        List of REE.
+    """
+    elements = [
+        "La",
+        "Ce",
+        "Pr",
+        "Nd",
+        "Pm",
+        "Sm",
+        "Eu",
+        "Gd",
+        "Tb",
+        "Dy",
+        "Ho",
+        "Er",
+        "Tm",
+        "Yb",
+        "Lu",
+    ]
+    if dropPm:
+        elements = [i for i in elements if not i == "Pm"]
+    if output == "formula":
+        elements = [getattr(pt, el) for el in elements]
+    return elements
+
+
+def REY(output="string", dropPm=True):
+    """
+    Provides a list of Rare Earth Elements, with the addition of Yttrium.
+
+    Parameters
+    -----------
+    output : :class:`str`
+        Whether to return output list as formulae ('formula') or strings (anthing else).
+
+    Returns
+    -------
+    :class:`list` | :class:`set`
+        List of REE+Y.
+
+    Notes
+    ------
+    This currently modifies the hardcoded list of :func:`REE`, but could be adapated
+    for different element ordering.
+    """
+    elements = [
+        "La",
+        "Ce",
+        "Pr",
+        "Nd",
+        "Pm",
+        "Sm",
+        "Eu",
+        "Gd",
+        "Tb",
+        "Dy",
+        "Y",
+        "Ho",
+        "Er",
+        "Tm",
+        "Yb",
+        "Lu",
+    ]
+    if dropPm:
+        elements = [i for i in elements if not i == "Pm"]
+    if output == "formula":
+        elements = [getattr(pt, el) for el in elements]
+    return elements
+
+
 def REE(output="string", dropPm=True):
     """
     Provides a list of Rare Earth Elements.

test/geochem/geochem_ind.py

@@ -108,6 +108,42 @@ def test_string_output(self):
                 self.assertIs(type(el), str)
 
 
+class TestREY(unittest.TestCase):
+    """Tests the Rare Earth Element + Yttrium generator."""
+
+    def setUp(self):
+        self.min_z = 57
+        self.max_z = 71
+
+    def test_complete(self):
+        """Check all REE and Yttrium are present."""
+        reels = REY(output="formula", dropPm=False)
+        ns = [el.number for el in reels]
+        print(ns)
+        for n in [39] + [*range(self.min_z, self.max_z + 1)]:
+            with self.subTest(n=n):
+                self.assertTrue(n in ns)
+
+    def test_precise(self):
+        """Check that only the REY are returned."""
+        reels = REY(output="formula")
+        ns = [el.number for el in reels]
+        self.assertTrue(min(ns) == 39)
+        self.assertTrue(max(ns) == self.max_z)
+
+    def test_formula_output(self):
+        """Check the function produces formula output."""
+        for el in REY(output="formula"):
+            with self.subTest(el=el):
+                self.assertIs(type(el), type(pt.elements[0]))
+
+    def test_string_output(self):
+        """Check the function produces string output."""
+        for el in REY(output="string"):
+            with self.subTest(el=el):
+                self.assertIs(type(el), str)
+
+
 class TestSimpleOxides(unittest.TestCase):
     """Tests the simple oxide generator."""
 

test/geochem/geochem_pyrochem.py

@@ -48,14 +48,14 @@ def test_pyrochem_indexes(self):
 
     def test_pyrochem_subsetters(self):
         obj = self.df
-        for subset in ["REE", "elements", "oxides", "isotope_ratios"]:
+        for subset in ["REE", "REY", "elements", "oxides", "isotope_ratios"]:
             with self.subTest(subset=subset):
                 out = getattr(obj.pyrochem, subset)
                 self.assertIsInstance(out, obj.__class__)  # in this case a dataframe
 
     def test_pyrochem_subsetter_assignment(self):
         obj = self.df
-        for subset in ["REE", "elements", "oxides", "isotope_ratios"]:
+        for subset in ["REE", "REY", "elements", "oxides", "isotope_ratios"]:
             with self.subTest(subset=subset):
                 setattr(obj.pyrochem, subset, getattr(obj.pyrochem, subset) * 1.0)