Repository: pynucastro/pynucastro

Issue:
make functions more robust to nuclei inputs
Right now, some functions take a string like `"c12"` and others require `Nucleus("c12")`,  We should make all functions that take a nucleus or list of nuclei work off of either strings or Nucleus objects.

For using in Jupyter, it is often easier to just pass in strings.

Discussion:

Patch:
diff --git a/pynucastro/networks/rate_collection.py b/pynucastro/networks/rate_collection.py
index c33e5923..b65cfa35 100644
--- a/pynucastro/networks/rate_collection.py
+++ b/pynucastro/networks/rate_collection.py
@@ -92,10 +92,11 @@ class Composition:
 
     """
     def __init__(self, nuclei, small=1.e-16):
-        """nuclei is an iterable of the nuclei (Nucleus objects) in the network"""
-        if not isinstance(nuclei[0], Nucleus):
-            raise ValueError("must supply an iterable of Nucleus objects")
-        self.X = {k: small for k in nuclei}
+        """nuclei is an iterable of the nuclei in the network"""
+        try:
+            self.X = {Nucleus.cast(k): small for k in nuclei}
+        except TypeError:
+            raise ValueError("must supply an iterable of Nucleus objects or strings") from None
 
     def __str__(self):
         ostr = ""
@@ -138,10 +139,8 @@ class Composition:
 
     def set_nuc(self, name, xval):
         """ set nuclei name to the mass fraction xval """
-        for k in self.X:
-            if k.raw == name:
-                self.X[k] = xval
-                break
+        nuc = Nucleus.cast(name)
+        self.X[nuc] = xval
 
     def normalize(self):
         """ normalize the mass fractions to sum to 1 """
@@ -177,6 +176,9 @@ class Composition:
         matches will be binned into the nuclei in that list
         """
 
+        nuclei = Nucleus.cast_list(nuclei)
+        exclude = Nucleus.cast_list(exclude, allow_None=True)
+
         # sort the input nuclei by A, then Z
         nuclei.sort(key=lambda n: (n.A, n.Z))
 
@@ -371,9 +373,9 @@ class RateCollection:
         If rates is supplied, initialize a RateCollection using the
         Rate objects in the list 'rates'.
 
-        inert_nuclei is a list of nuclei (as Nucleus objects) that
-        should be part of the collection but are not linked via reactions
-        to the other nuclei in the network.
+        inert_nuclei is a list of nuclei that should be part of the
+        collection but are not linked via reactions to the other nuclei
+        in the network.
 
         symmetric_screening means that we screen the reverse rates
         using the same factor as the forward rates, for rates computed
@@ -387,13 +389,11 @@ class RateCollection:
         self.rates = []
         combined_library = Library()
 
-        self.inert_nuclei = inert_nuclei
+        self.inert_nuclei = Nucleus.cast_list(inert_nuclei, allow_None=True)
 
         self.symmetric_screening = symmetric_screening
         self.do_screening = do_screening
 
-        self.inert_nuclei = inert_nuclei
-
         if rate_files:
             if isinstance(rate_files, str):
                 rate_files = [rate_files]
@@ -438,12 +438,8 @@ class RateCollection:
                 if r.intermediate_nucleus not in self.unique_nuclei + self.approx_nuclei:
                     self.approx_nuclei.append(r.intermediate_nucleus)
 
-        if self.inert_nuclei:
-            for n in self.inert_nuclei:
-                if isinstance(n, Nucleus):
-                    nuc = n
-                else:
-                    nuc = Nucleus(n)
+        if self.inert_nuclei is not None:
+            for nuc in self.inert_nuclei:
                 if nuc not in self.unique_nuclei:
                     self.unique_nuclei.append(nuc)
 
@@ -649,9 +645,11 @@ class RateCollection:
 
     def get_rate_by_nuclei(self, reactants, products):
         """given a list of reactants and products, return any matching rates"""
+        reactants = sorted(Nucleus.cast_list(reactants))
+        products = sorted(Nucleus.cast_list(products))
         _tmp = [r for r in self.rates if
-                sorted(r.reactants) == sorted(reactants) and
-                sorted(r.products) == sorted(products)]
+                sorted(r.reactants) == reactants and
+                sorted(r.products) == products]
 
         if not _tmp:
             return None
@@ -708,13 +706,11 @@ class RateCollection:
         that directly involve them (this doesn't affect approximate rates that
         may have these nuclei as hidden intermediate links)"""
 
+        nuc_list = Nucleus.cast_list(nuc_list)
         rates_to_delete = []
         for nuc in nuc_list:
-            nn = nuc
-            if not isinstance(nuc, Nucleus):
-                nn = Nucleus(nuc)
             for rate in self.rates:
-                if nn in rate.reactants + rate.products:
+                if nuc in rate.reactants + rate.products:
                     print(f"looking to remove {rate}")
                     rates_to_delete.append(rate)
 
@@ -755,13 +751,7 @@ class RateCollection:
         effective approximate rate."""
 
         # make sure that the intermediate_nuclei list are Nuclei objects
-        _inter_nuclei_remove = []
-        if intermediate_nuclei is not None:
-            for nn in intermediate_nuclei:
-                if isinstance(nn, Nucleus):
-                    _inter_nuclei_remove.append(nn)
-                else:
-                    _inter_nuclei_remove.append(Nucleus(nn))
+        intermediate_nuclei = Nucleus.cast_list(intermediate_nuclei, allow_None=True)
 
         # find all of the (a,g) rates
         ag_rates = []
@@ -784,7 +774,7 @@ class RateCollection:
 
             inter_nuc = Nucleus(f"{element.abbreviation}{inter_nuc_A}")
 
-            if intermediate_nuclei and inter_nuc not in _inter_nuclei_remove:
+            if intermediate_nuclei and inter_nuc not in intermediate_nuclei:
                 continue
 
             # look for A(a,p)X
diff --git a/pynucastro/nucdata/nucleus.py b/pynucastro/nucdata/nucleus.py
index 88bb3e52..5b9fc5be 100644
--- a/pynucastro/nucdata/nucleus.py
+++ b/pynucastro/nucdata/nucleus.py
@@ -183,6 +183,24 @@ class Nucleus:
             return self.Z < other.Z
         return self.A < other.A
 
+    @classmethod
+    def cast(cls, obj):
+        if isinstance(obj, cls):
+            return obj
+        if isinstance(obj, str):
+            return cls.from_cache(obj)
+        raise TypeError("Invalid type passed to Nucleus.cast() (expected str or Nucleus)")
+
+    @classmethod
+    def cast_list(cls, lst, *, allow_None=False, allow_single=False):
+        if allow_None and lst is None:
+            return lst
+        if isinstance(lst, (str, cls)):
+            if allow_single:
+                return [cls.cast(lst)]
+            raise ValueError("Single object passed to Nucleus.cast_list() instead of list")
+        return [cls.cast(obj) for obj in lst]
+
 
 def get_nuclei_in_range(zmin, zmax, amin, amax):
     """given a range of Z = [zmin, zmax], and A = [amin, amax],
diff --git a/pynucastro/rates/library.py b/pynucastro/rates/library.py
index b9565eb1..7cbf1f62 100644
--- a/pynucastro/rates/library.py
+++ b/pynucastro/rates/library.py
@@ -307,9 +307,11 @@ class Library:
 
     def get_rate_by_nuclei(self, reactants, products):
         """given a list of reactants and products, return any matching rates"""
+        reactants = sorted(Nucleus.cast_list(reactants))
+        products = sorted(Nucleus.cast_list(products))
         _tmp = [r for r in self.get_rates() if
-                sorted(r.reactants) == sorted(reactants) and
-                sorted(r.products) == sorted(products)]
+                sorted(r.reactants) == reactants and
+                sorted(r.products) == products]
 
         if not _tmp:
             return None
@@ -418,16 +420,7 @@ class Library:
         If print_warning is True, then print out a warning if one of the input nuclei is not linked.
         """
 
-        if isinstance(nuclist, (Nucleus, str)):
-            nuclist = [nuclist]
-
-        nucleus_set = set()
-        for nuc in nuclist:
-            if isinstance(nuc, Nucleus):
-                nucleus_set.add(nuc)
-            else:
-                anuc = Nucleus(nuc)
-                nucleus_set.add(anuc)
+        nucleus_set = set(Nucleus.cast_list(nuclist))
 
         # Discard rates with nuclei that are not in nucleus_set
         filtered_rates = []
@@ -617,21 +610,9 @@ class RateFilter:
         self.filter_function = filter_function
 
         if reactants:
-            if isinstance(reactants, (Nucleus, str)):
-                reactants = [reactants]
-            self.reactants = [self._cast_nucleus(r) for r in reactants]
+            self.reactants = Nucleus.cast_list(reactants, allow_single=True)
         if products:
-            if isinstance(products, (Nucleus, str)):
-                products = [products]
-            self.products = [self._cast_nucleus(r) for r in products]
-
-    @staticmethod
-    def _cast_nucleus(r):
-        """ Make sure r is of type Nucleus. """
-        if not isinstance(r, Nucleus):
-            rnuc = Nucleus(r)
-            return rnuc
-        return r
+            self.products = Nucleus.cast_list(products, allow_single=True)
 
     @staticmethod
     def _contents_equal(a, b):
diff --git a/pynucastro/rates/rate.py b/pynucastro/rates/rate.py
index a23f592d..4607ad22 100644
--- a/pynucastro/rates/rate.py
+++ b/pynucastro/rates/rate.py
@@ -334,12 +334,12 @@ class Rate:
         sources.  Here we only specify the reactants and products and Q value"""
 
         if reactants:
-            self.reactants = reactants
+            self.reactants = Nucleus.cast_list(reactants)
         else:
             self.reactants = []
 
         if products:
-            self.products = products
+            self.products = Nucleus.cast_list(products)
         else:
             self.products = []
 
@@ -803,12 +803,12 @@ class ReacLibRate(Rate):
         self.fname = None
 
         if reactants:
-            self.reactants = reactants
+            self.reactants = Nucleus.cast_list(reactants)
         else:
             self.reactants = []
 
         if products:
-            self.products = products
+            self.products = Nucleus.cast_list(products)
         else:
             self.products = []
 
@@ -888,16 +888,7 @@ class ReacLibRate(Rate):
             self.fname += "__removed"
 
     def modify_products(self, new_products):
-        if not isinstance(new_products, (set, list, tuple)):
-            new_products = [new_products]
-
-        self.products = []
-        for p in new_products:
-            if isinstance(p, Nucleus):
-                self.products.append(p)
-            else:
-                self.products.append(Nucleus(p))
-
+        self.products = Nucleus.cast_list(new_products, allow_single=True)
         self.modified = True
 
         # we need to update the Q value and the print string for the rate
diff --git a/pynucastro/reduction/drgep.py b/pynucastro/reduction/drgep.py
index 9ff54365..0bc7ea97 100644
--- a/pynucastro/reduction/drgep.py
+++ b/pynucastro/reduction/drgep.py
@@ -1,5 +1,6 @@
 import numpy as np
 
+from pynucastro.nucdata import Nucleus
 from pynucastro.reduction.reduction_utils import (mpi_importer,
                                                   mpi_numpy_decomp, to_list)
 
@@ -297,7 +298,7 @@ def drgep(net, conds, targets, tols, returnobj='net', use_mpi=False, use_numpy=F
         latter case, the sequences will be permuted to create the dataset. The compositions should
         be pynucastro Composition objects.
     :param targets: A collection of target nuclei (or a single target nucleus) to run the
-        graph search algorithm from. Should be supplied as pynucastro Nucleus objects.
+        graph search algorithm from. Should be supplied as pynucastro Nucleus objects or strings.
     :param tols: Tolerance(s) or cutoff threshold(s) to use for paths from each of the target nuclei.
         Nuclei whose interaction coefficients do not meet the specified tolerance will have their
         interaction coefficients set to 0.0. Can be a single number (will be the same for all targets)
@@ -320,7 +321,7 @@ def drgep(net, conds, targets, tols, returnobj='net', use_mpi=False, use_numpy=F
 
     if returnobj not in {'net', 'nuclei', 'coeff'}:
         raise ValueError(f"Invalid 'returnobj' argument: '{returnobj}'.")
-    targets = to_list(targets)
+    targets = Nucleus.cast_list(targets)
     tols = to_list(tols, len(targets))
 
     #-------------------------------------------------------
diff --git a/pynucastro/reduction/load_network.py b/pynucastro/reduction/load_network.py
index 0c520b33..e28b2a7b 100755
--- a/pynucastro/reduction/load_network.py
+++ b/pynucastro/reduction/load_network.py
@@ -13,6 +13,8 @@ def load_network(endpoint=Nucleus('te108'), library_name='rp-process-lib'):
     the endpoint's Z and A are included.
     """
 
+    endpoint = Nucleus.cast(endpoint)
+
     def ff(rate):
 
         react_meet_conds = \
diff --git a/pynucastro/reduction/sensitivity_analysis.py b/pynucastro/reduction/sensitivity_analysis.py
index e6cf7c87..1870386f 100644
--- a/pynucastro/reduction/sensitivity_analysis.py
+++ b/pynucastro/reduction/sensitivity_analysis.py
@@ -1,5 +1,6 @@
 import numpy as np
 
+from pynucastro.nucdata import Nucleus
 from pynucastro.reduction.reduction_utils import mpi_importer
 
 MPI = mpi_importer()
@@ -25,6 +26,8 @@ def binary_search_trim(network, nuclei, errfunc, thresh=0.05):
         threshold.
     """
 
+    nuclei = Nucleus.cast_list(nuclei)
+
     start_idx = 0
     seg_size = len(nuclei) / 2
 
diff --git a/pynucastro/screening/screen.py b/pynucastro/screening/screen.py
index aedd74f2..68a75d4b 100644
--- a/pynucastro/screening/screen.py
+++ b/pynucastro/screening/screen.py
@@ -4,6 +4,7 @@ Python implementations of screening routines.
 import numpy as np
 
 from pynucastro.constants import constants
+from pynucastro.nucdata import Nucleus
 # use the jitclass placeholder from rate.py
 from pynucastro.rates.rate import jitclass, numba
 
@@ -155,6 +156,8 @@ def make_screen_factors(n1, n2):
     :param Nucleus n1: first nucleus
     :param Nucleus n2: second nucleus
     """
+    n1 = Nucleus.cast(n1)
+    n2 = Nucleus.cast(n2)
     return ScreenFactors(n1.Z, n1.A, n2.Z, n2.A)
 
 


Test patch:
diff --git a/pynucastro/nucdata/tests/test_nucleus.py b/pynucastro/nucdata/tests/test_nucleus.py
index ecfd31dc..6d350680 100644
--- a/pynucastro/nucdata/tests/test_nucleus.py
+++ b/pynucastro/nucdata/tests/test_nucleus.py
@@ -1,4 +1,4 @@
-from pytest import approx
+from pytest import approx, raises
 
 from pynucastro.nucdata import Nucleus, get_nuclei_in_range
 
@@ -111,3 +111,32 @@ class TestNucleus:
         assert nuc_list[12] == Nucleus("o14")
         assert nuc_list[13] == Nucleus("o15")
         assert nuc_list[14] == Nucleus("o16")
+
+    def test_cast(self):
+        assert Nucleus.cast("c12") == self.c12
+        assert Nucleus.cast("C12") == self.c12
+        assert Nucleus.cast(self.c12) == self.c12
+        assert Nucleus.cast("n") == Nucleus("n")
+
+    def test_cast_list(self):
+        expected = [self.p, self.n, self.he4, self.c12, self.pb237]
+        assert Nucleus.cast_list(["p", self.n, "a", "c12", self.pb237]) == expected
+        assert Nucleus.cast_list(["p", "n", "he4", "c12", "pb237"]) == expected
+        assert Nucleus.cast_list(expected) == expected
+        assert Nucleus.cast_list([self.p, "n", "he4", self.c12, self.pb237], allow_single=True) == expected
+        assert Nucleus.cast_list(expected, allow_None=True, allow_single=True) == expected
+        assert Nucleus.cast_list(expected, allow_None=True, allow_single=True) == expected
+
+        with raises(ValueError):
+            Nucleus.cast_list("he4")
+        with raises(ValueError):
+            Nucleus.cast_list(self.he4)
+        assert Nucleus.cast_list(self.he4, allow_single=True) == [self.he4]
+        assert Nucleus.cast_list("he4", allow_single=True) == [self.he4]
+        assert Nucleus.cast_list([self.he4], allow_single=True) == [self.he4]
+        assert Nucleus.cast_list(["he4"]) == [self.he4]
+
+        with raises(TypeError):
+            Nucleus.cast_list(None)
+        assert Nucleus.cast_list(None, allow_None=True) is None
+        assert Nucleus.cast_list([]) == []