Lint: Formats/adds type annotations to kTP code

autoreact/__init__.py

@@ -1,10 +1,10 @@
-""" autoreact
+"""autoreact.
 
 I this should be renamed `autorate` and house the functions in ratefit and
 mechanalyzer.calculator.rates, along with the parametrization functions.
 """
 from autoreact import params
 
 __all__ = [
-    'params',
+    "params",
 ]

autoreact/ktp_xarray/xarray_wrappers.py

@@ -1,85 +1,85 @@
-"""
-Wrappers for the new xarray system. Constructors, Getters, then Setters.
-"""
+"""Wrappers for the new xarray system. Constructors, Getters, then Setters."""
+
+from collections.abc import Sequence
+from numbers import Number
 
-import xarray
 import numpy
+import xarray
 
-# Constructors
-def from_data(temps, press, rates):
-    """Construct a KTP DataArray from data"""
+RateConstant = xarray.DataArray
 
-    ktp = xarray.DataArray(rates, (("pres", press), ("temp", temps)))
 
+# Constructors
+def from_data(
+    temps: Sequence[Number], press: Sequence[Number], rates: Sequence[Sequence[Number]]
+) -> RateConstant:
+    """Construct a KTP DataArray from data.."""
+    data = [list(map(float, kT)) for kT in rates]
+    coords = {
+        "pres": list(map(float, press)),
+        "temp": list(map(float, temps)),
+    }
+    ktp = xarray.DataArray(data=data, coords=coords)
     return ktp
 
 
-
 # Getters
-def get_pressures(ktp):
-    """Gets the pressure values"""
-
+def get_pressures(ktp: RateConstant) -> numpy.ndarray:
+    """Get the pressure values.."""
     return ktp.pres.data
 
 
-def get_temperatures(ktp):
-    """Gets the temperature values"""
-
+def get_temperatures(ktp: RateConstant) -> numpy.ndarray:
+    """Get the temperature values.."""
     return ktp.temp.data
 
 
-def get_values(ktp):
-    """Gets the KTP values"""
-
+def get_values(ktp: RateConstant) -> numpy.ndarray:
+    """Get the KTP values.."""
     return ktp.values
 
 
-def get_pslice(ktp, ip):
-    """Get a slice at a selected pressure value"""
+def get_pslice(ktp: RateConstant, p: Number) -> numpy.ndarray:
+    """Get a slice at a selected pressure value."""
+    return ktp.sel(pres=p)
 
-    return ktp.sel(pres=ip)
 
+def get_tslice(ktp: RateConstant, t: Number) -> numpy.ndarray:
+    """Get a slice at a selected temperature value."""
+    return ktp.sel(temp=t)
 
-def get_tslice(ktp, it):
-    """Get a slice at a selected temperature value"""
 
-    return ktp.sel(temp=it)
+def get_spec_vals(ktp: RateConstant, t: Number, p: Number) -> numpy.ndarray:
+    """Get a specific value at a selected temperature and pressure value."""
+    return ktp.sel(temp=t, pres=p)
 
 
-def get_spec_vals(ktp, it, ip):
-    """Get a specific value at a selected temperature and pressure value"""
-
-    return ktp.sel(temp=it, pres=ip)
-
-
-def get_ipslice(ktp, ip):
-    """Get a slice at a selected pressure index"""
-
+def get_ipslice(ktp: RateConstant, ip: int) -> numpy.ndarray:
+    """Get a slice at a selected pressure index."""
     return ktp.isel(pres=ip)
 
 
-def get_itslice(ktp, it):
-    """Get a slice at a selected temperature index"""
-
+def get_itslice(ktp: RateConstant, it: int):
+    """Get a slice at a selected temperature index."""
     return ktp.isel(temp=it)
 
 
-
 # Setters
-def set_rates(ktp, rates, pres, temp):
-    """Sets the KTP values"""
-
+def set_rates(
+    ktp: RateConstant, rates: Sequence[Sequence[Number]], pres: Number, temp: Number
+):
+    """Set the KTP values."""
     ktp.loc[{"pres": pres, "temp": temp}] = rates
     return ktp
 
 
 # Translators
 
-def dict_from_xarray(xarray_in):
-    """Turns an xarray into a ktp_dct"""
 
+def dict_from_xarray(xarray_in: RateConstant) -> dict:
+    """Turn an xarray into a ktp_dct."""
     ktp_dct = {}
-    #dict_temps = get_temperatures(xarray)
+    # dict_temps = get_temperatures(xarray)
     for pres in get_pressures(xarray_in):
         dict_temps = get_temperatures(xarray_in)
         dict_kts = []
@@ -93,15 +93,16 @@ def dict_from_xarray(xarray_in):
         dict_temps = curr_temps
         dict_kts = numpy.array(dict_kts, dtype=numpy.float64)
         if pres == numpy.inf:
-            pres = 'high'
+            pres = "high"
         ktp_dct[pres] = (dict_temps, dict_kts)
     return ktp_dct
 
-#Stopped working on this one because it was less critical!
 
-#def xarray_from_dict(ktp_dct):
+# Stopped working on this one because it was less critical!
+
+# def xarray_from_dict(ktp_dct):
 #    """DOES NOT WORK YET!
-#    Turns a ktp_dct into an xarray"""
+#    Turn a ktp_dct into an xarray."""
 #
 #    xarray_press = []
 #    xarray_temps = []

autoreact/ktp_xarray/xarray_wrappers_test.py

@@ -1,93 +1,99 @@
-"""Tests xarray_wrappers.py's functions"""
+"""Test xarray_wrappers.py's functions."""
 
 import numpy
+
 from autoreact.ktp_xarray import xarray_wrappers
 
 Temps = [1000, 1500, 2000, 2500]
 Press = [1, 10, numpy.inf]
 Rates = [[1e1, 1e2, 1e3, 1e4], [1e5, 1e6, 1e7, 1e8], [1e9, 1e10, 1e11, 1e12]]
 
 Ktp = xarray_wrappers.from_data(Temps, Press, Rates)
-Ktp_dct = {1.0: (([1000., 1500., 2000., 2500.]),
-                ([ 10., 100., 1000., 10000.])),
-           10: (([1000., 1500., 2000., 2500.]),
-                ([1.e+05, 1.e+06, 1.e+07, 1.e+08])),
-           numpy.inf: (([1000., 1500., 2000., 2500.]),
-                ([1.e+09, 1.e+10, 1.e+11, 1.e+12]))}
+Ktp_dct = {
+    1.0: (([1000.0, 1500.0, 2000.0, 2500.0]), ([10.0, 100.0, 1000.0, 10000.0])),
+    10: (([1000.0, 1500.0, 2000.0, 2500.0]), ([1.0e05, 1.0e06, 1.0e07, 1.0e08])),
+    numpy.inf: (([1000.0, 1500.0, 2000.0, 2500.0]), ([1.0e09, 1.0e10, 1.0e11, 1.0e12])),
+}
 print(Ktp)
 
+
 def test_get_temperatures():
-    """Tests the get_temperatures function"""
+    """Test the get_temperatures function."""
     temp = xarray_wrappers.get_temperatures(Ktp)
     print(temp)
 
 
 def test_get_pressures():
-    """Tests the get_pressures function"""
+    """Test the get_pressures function."""
     pres = xarray_wrappers.get_pressures(Ktp)
-    print(pres)
+    print(f"pres = {pres}")
+    print(type(pres))
 
 
 def test_get_values():
-    """Tests the get_values function"""
+    """Test the get_values function."""
     vals = xarray_wrappers.get_values(Ktp)
-    print(vals)
+    print(f"vals = {vals}")
+    print(type(vals))
 
 
 def test_get_pslice():
-    """Tests the get_pslice function"""
+    """Test the get_pslice function."""
     pslice = xarray_wrappers.get_pslice(Ktp, numpy.inf)
-    print(pslice)
+    print(f"pslice = {pslice}")
+    print(type(pslice))
 
 
 def test_get_tslice():
-    """Tests the get_tslice function"""
+    """Test the get_tslice function."""
     tslice = xarray_wrappers.get_tslice(Ktp, 1500)
     print(tslice)
 
 
 def test_get_spec_vals():
-    """Tests the get_spec_values function"""
+    """Test the get_spec_values function."""
     vals = xarray_wrappers.get_spec_vals(Ktp, 1500, 1)
     print(vals)
 
 
 def test_get_ipslice():
-    """Tests the get_ipslice function"""
+    """Test the get_ipslice function."""
     ipslice = xarray_wrappers.get_ipslice(Ktp, 0)
     print(ipslice)
 
 
 def test_get_itslice():
-    """Tests the get_itslice function"""
+    """Test the get_itslice function."""
     itslice = xarray_wrappers.get_itslice(Ktp, 0)
     print(itslice)
 
 
 def test_set_rates():
-    """Tests the set_rates function"""
+    """Test the set_rates function."""
     new_rates = xarray_wrappers.set_rates(Ktp, numpy.nan, 10, 2000)
     print(new_rates)
 
+
 def test_dict_from_xarray():
-    """Tests the ktp_to_xarray function"""
+    """Test the ktp_to_xarray function."""
     ktp_dct = xarray_wrappers.dict_from_xarray(Ktp)
     print(ktp_dct)
 
-#def test_xarray_from_dict():
-#    """Tests the set_ktp_dct function"""
+
+# def test_xarray_from_dict():
+#    """Test the set_ktp_dct function."""
 #    xarray = xarray_wrappers.xarray_from_dict(Ktp_dct)
 #    print(xarray)
 
 
-test_get_pressures()
-test_get_temperatures()
-test_get_values()
+# test_get_pressures()
+# test_get_temperatures()
+# test_get_values()
 test_get_pslice()
-test_get_tslice()
-test_get_spec_vals()
-test_get_ipslice()
-test_get_itslice()
-test_set_rates()
-test_dict_from_xarray()
-#test_xarray_from_dict()
+# test_get_tslice()
+# test_get_spec_vals()
+# test_get_ipslice()
+# test_get_itslice()
+# test_set_rates()
+# test_dict_from_xarray()
+# test_xarray_from_dict()