From 371d0ff1f7913826ce4a0592668a6b3319e7280d Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Thu, 30 Apr 2020 15:43:32 -0400 Subject: [PATCH 01/49] Include HDF5 files in MANIFEST.in --- MANIFEST.in | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/MANIFEST.in b/MANIFEST.in index fb17c8c..e1fc8c1 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -4,7 +4,7 @@ include setup.cfg include LICENSE.md include pyproject.toml -recursive-include plasmapy_nei *.pyx *.c *.pxd +recursive-include plasmapy_nei *.pyx *.c *.pxd *.h5 recursive-include docs * recursive-include licenses * recursive-include cextern * From d707616e13a4608ebe172adfaf3e4069b67d54a1 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Thu, 30 Apr 2020 15:46:26 -0400 Subject: [PATCH 02/49] Transfer EigenData2 class and tests from NEI-modeling/NEI Co-authored-by: Chengcai Shen --- plasmapy_nei/eigen/eigenvaluetable.py | 342 ++++++++++++++++++ .../eigen/tests/test_eigenvaluetable.py | 180 +++++++++ 2 files changed, 522 insertions(+) create mode 100644 plasmapy_nei/eigen/eigenvaluetable.py create mode 100644 plasmapy_nei/eigen/tests/test_eigenvaluetable.py diff --git a/plasmapy_nei/eigen/eigenvaluetable.py b/plasmapy_nei/eigen/eigenvaluetable.py new file mode 100644 index 0000000..3cad9d5 --- /dev/null +++ b/plasmapy_nei/eigen/eigenvaluetable.py @@ -0,0 +1,342 @@ +"""The EigenData2 class.""" + +import warnings +import numpy as np +from numpy import linalg as LA +from plasmapy import atomic +from .. import __path__ +import h5py + + +class EigenData2: + """ + + A class to obtain eigenvalue tables used in NEI calculations. + + Parameters + ---------- + element : `str` or `int` + A string representing a element symbol. The defaut value + is for Hydrogen. + + Raises: + ---------- + + Properties: + temperature: `float` + The temperary is set to get the relative rates and eigen values. + In unit of K. + + Examples: + ---------- + To get the table for element 'Helium' at Te=5.0e5K: + + >>> table = EigenData2(element=2) + >>> table.temperature=5.0e5 + + Output eigenvals: + >>> table.eigenvalues() + array([-1.12343827e-08, -9.00005970e-10, 8.58945565e-30]) + + Output equilibrium states: + >>> table.equilibrium_state() + array([9.57162006e-09, 1.26564673e-04, 9.99873426e-01]) + + Or one may output properties at other temperatures, for example: + >>> table.eigenvalues(T_e=12589.254117941662) + array([-1.87480784e-12, -3.73191378e-13, 0.00000000e+00]) + + Or input temperature index on the grid: + >>> table.eigenvalues(T_e_index=100) + array([-1.87480784e-12, -3.73191378e-13, 0.00000000e+00]) + + """ + + def __init__(self, element='H'): + """Read in the """ + + self._element = element + self._temperature = None + + # + # 1. Read ionization and recombination rates + # + data_dir = __path__[0] + '/data/ionizrecombrates/chianti_8.07/' + filename = data_dir + 'ionrecomb_rate.h5' + f = h5py.File(filename, 'r') + + atomic_numb = atomic.atomic_number(element) + nstates = atomic_numb + 1 + + self._temperature_grid = f['te_gird'][:] + ntemp = len(self._temperature_grid) + c_ori = f['ioniz_rate'][:] + r_ori = f['recomb_rate'][:] + f.close() + + # + # Ionization and recombination rate for the current element + # + c_rate = np.zeros((ntemp, nstates)) + r_rate = np.zeros((ntemp, nstates)) + for ite in range(ntemp): + for i in range(nstates-1): + c_rate[ite, i] = c_ori[i, atomic_numb-1, ite] + for i in range(1, nstates): + r_rate[ite, i] = r_ori[i-1, atomic_numb-1, ite] + + # + # 2. Definet the grid size + # + self._ntemp = ntemp + self._atomic_numb = atomic_numb + self._nstates = nstates + + # + # Compute eigenvalues and eigenvectors + # + self._ionization_rate = np.ndarray(shape=(ntemp, nstates), + dtype=np.float64) + + self._recombination_rate = np.ndarray(shape=(ntemp, nstates), + dtype=np.float64) + + self._equilibrium_states = np.ndarray(shape=(ntemp, nstates), + dtype=np.float64) + + self._eigenvalues = np.ndarray(shape=(ntemp, nstates), + dtype=np.float64) + + self._eigenvectors = np.ndarray(shape=(ntemp, nstates, nstates), + dtype=np.float64) + + self._eigenvector_inverses = np.ndarray( + shape=(ntemp, nstates, nstates), + dtype=np.float64) + + # + # Save ionization and recombination rates + # + self._ionization_rate = c_rate + self._recombination_rate = r_rate + + # + # Define the coefficients matrix A. The first dimension is + # for elements, and the second number of equations. + # + neqs = nstates + A = np.ndarray(shape=(nstates, neqs), dtype=np.float64) + + # + # Enter temperature loop over the whole temperature grid + # + for ite in range(ntemp): + # Ionization and recombination rate at Te(ite) + carr = c_rate[ite, :] + rarr = r_rate[ite, :] + + # Equilibirum + eqi = self._function_eqi(carr, rarr, atomic_numb) + + # Initialize A to zero + for ion in range(nstates): + for jon in range(nstates): + A[ion, jon] = 0.0 + + # Give coefficients + for ion in range(1, nstates-1): + A[ion, ion-1] = carr[ion-1] + A[ion, ion] = -(carr[ion]+rarr[ion]) + A[ion, ion+1] = rarr[ion+1] + + # The first row + A[0, 0] = -carr[0] + A[0, 1] = rarr[1] + + # The last row + A[nstates-1, nstates-2] = carr[nstates-2] + A[nstates-1, nstates-1] = -rarr[nstates-1] + + # Compute eigenvalues and eigenvectors using Scipy + la, v = LA.eig(A) + + # Rerange the eigenvalues. Try a simple way in here. + idx = np.argsort(la) + la = la[idx] + v = v[:, idx] + + # Compute inverse of eigenvectors + v_inverse = LA.inv(v) + + # transpose the order to as same as the Fortran Version + v = v.transpose() + v_inverse = v_inverse.transpose() + + # Save eigenvalues and eigenvectors into arrays + for j in range(nstates): + self._eigenvalues[ite, j] = la[j] + self._equilibrium_states[ite, j] = eqi[j] + for i in range(nstates): + self._eigenvectors[ite, i, j] = v[i, j] + self._eigenvector_inverses[ite, i, j] = v_inverse[i, j] + + # + # The following Functions is used to obtain the eigen values and relative + # def properties. + # + def _get_temperature_index(self, T_e): + """Returns the temperature index closest to a particular + temperature.""" + T_e_array = self._temperature_grid + + # Check the temperature range + T_e_grid_max = np.amax(T_e_array) + T_e_grid_min = np.amin(T_e_array) + + if T_e == T_e_grid_max: + return self._ntemp - 1 + elif T_e == T_e_grid_min: + return 0 + elif T_e > T_e_grid_max: + warnings.warn(f"Temperature exceeds the temperature grid " + f"boundary: temperature index will be reset " + f"to {self._ntemp - 1}", UserWarning) + return self._ntemp-1 + elif T_e < T_e_grid_min: + warnings.warn(f"Temperature is below the temperature grid " + f"boundary: temperature index will be reset to " + f"0.", UserWarning) + return 0 + + # TODO: Add a test to check that the temperature grid is monotonic + + res = np.where(T_e_array >= T_e) + res_ind = res[0] + index = res_ind[0] + dte_l = abs(T_e - T_e_array[index - 1]) # re-check the neighbor point + dte_r = abs(T_e - T_e_array[index]) + if (dte_l <= dte_r): + index = index - 1 + return index + + @property + def temperature(self): + """Returns the electron temperature currently in use by this class, + or None if the temperature has not been set.""" + return self._temperature + + @temperature.setter + def temperature(self, T_e): + """Sets the electron temperature and index on the temperature grid + to be used by this class""" + # TODO: Add checks for the temperature + self._temperature = T_e + self._te_index = self._get_temperature_index(T_e) + + @property + def temperature_grid(self): + """Returns the grid of temperatures corresponding to the eigendata.""" + return self._temperature_grid + + def eigenvalues(self, T_e=None, T_e_index=None): + """Returns the eigenvalues for the ionization and recombination + rates for the temperature specified in the class.""" + if T_e_index: + return self._eigenvalues[T_e_index, :] + elif T_e: + T_e_index = self._get_temperature_index(T_e) + return self._eigenvalues[T_e_index, :] + elif self.temperature: + return self._eigenvalues[self._te_index, :] + else: + raise AttributeError("The temperature has not been set.") + + def eigenvectors(self, T_e=None, T_e_index=None): + """Returns the eigenvectors for the ionization and recombination + rates for the temperature specified in the class.""" + if T_e_index: + return self._eigenvectors[T_e_index, :, :] + elif T_e: + T_e_index = self._get_temperature_index(T_e) + return self._eigenvectors[T_e_index, :, :] + elif self.temperature: + return self._eigenvectors[self._te_index, :, :] + else: + raise AttributeError("The temperature has not been set.") + + def eigenvector_inverses(self, T_e=None, T_e_index=None): + """Returns the inverses of the eigenvectors for the ionization and + recombination rates for the temperature specified in the class.""" + if T_e_index: + return self._eigenvector_inverses[T_e_index, :, :] + elif T_e: + T_e_index = self._get_temperature_index(T_e) + return self._eigenvector_inverses[T_e_index, :, :] + elif self.temperature: + return self._eigenvector_inverses[self._te_index, :, :] + else: + raise AttributeError("The temperature has not been set.") + + + def equilibrium_state(self, T_e=None, T_e_index=None): + """Returns the equilibrium charge state distribution for the + temperature specified in the class.""" + if T_e_index: + return self._equilibrium_states[T_e_index, :] + elif T_e: + T_e_index = self._get_temperature_index(T_e) + return self._equilibrium_states[T_e_index, :] + elif self.temperature: + return self._equilibrium_states[self._te_index, :] + else: + raise AttributeError("The temperature has not been set.") + + def _function_eqi(self, ioniz_rate, recomb_rate, natom): + """Compute the equilibrium charge state distribution for the + temperature specified using ionization and recombinaiton rates.""" + nstates = natom + 1 + conce = np.zeros(nstates) + f = np.zeros(nstates + 1) + c = np.zeros(nstates + 1) + r = np.zeros(nstates + 1) + + # The start index is 1. + for i in range(nstates): + c[i+1] = ioniz_rate[i] + r[i+1] = recomb_rate[i] + + # Set f1 + f[1] = 1.0 + # f2 = c1*f1/r2 + f[2] = c[1]*f[1]/r[2] + # + # For Hydrogen, the following loop is not necessary. + # + if (natom <= 1): + f[1] = 1.0/(1.0 + c[1]/r[2]) + f[2] = c[1]*f[1]/r[2] + conce[0:2] = f[1:3] + return conce + + # + # For other elements with atomic number >= 2: + # + # f(i+1) = -(c(i-1)*f(i-1) - (c(i)+r(i)*f(i)))/r(i+1) + for k in range(2, natom): + f[k+1] = (-c[k-1]*f[k-1] + (c[k]+r[k])*f[k])/r[k+1] + + # f(natom+1) = c(natom)*f(natom)/r(natom+1) + f[natom+1] = c[natom]*f[natom]/r[natom+1] + # f1 = 1/sum(f(*)) + f[1] = 1.0/np.sum(f) + # f2 = c1*f1/r2 + f[2] = c[1]*f[1]/r[2] + # f(i+1) = -(c(i-1)*f(i-1) - (c(i)+r(i)*f(i)))/r(i+1) + for k in range(2, natom): + f[k+1] = (-c[k-1]*f[k-1] + (c[k]+r[k])*f[k])/r[k+1] + + # f(natom+1) = c(natom)*f(natom)/r(natom+1) + f[natom+1] = c[natom]*f[natom]/r[natom+1] + # + conce[0:nstates] = f[1:nstates+1] + return conce diff --git a/plasmapy_nei/eigen/tests/test_eigenvaluetable.py b/plasmapy_nei/eigen/tests/test_eigenvaluetable.py new file mode 100644 index 0000000..fc34396 --- /dev/null +++ b/plasmapy_nei/eigen/tests/test_eigenvaluetable.py @@ -0,0 +1,180 @@ +"""Test_eigenvaluetable""" +import warnings +import numpy as np +from plasmapy import atomic +from ..eigenvaluetable import EigenData2 +import pytest + +#------------------------------------------------------------------------------- +# Set elements and temperary list as testing inputs +#------------------------------------------------------------------------------- +natom_list = np.arange(1, 27) +natom = 8 + +#------------------------------------------------------------------------------ +# function: Time-Advance solover +#------------------------------------------------------------------------------ +def func_solver_eigenval(natom, te, ne, dt, f0, table): + """ + The testing function for performing time_advance calculations. + """ + + common_index = table._get_temperature_index(te) + evals = table.eigenvalues(T_e_index=common_index) # find eigenvalues on the chosen Te node + evect = table.eigenvectors(T_e_index=common_index) + evect_invers = table.eigenvector_inverses(T_e_index=common_index) + + # define the temperary diagonal matrix + diagona_evals = np.zeros((natom+1, natom+1)) + for ii in range(0, natom+1): + diagona_evals[ii,ii] = np.exp(evals[ii]*dt*ne) + + # matirx operation + matrix_1 = np.dot(diagona_evals, evect) + matrix_2 = np.dot(evect_invers, matrix_1) + + # get ions fraction at (time+dt) + ft = np.dot(f0, matrix_2) + + # re-check the smallest value + minconce = 1.0e-15 + for ii in np.arange(0, natom+1, dtype=np.int): + if (abs(ft[ii]) <= minconce): + ft[ii] = 0.0 + return ft + +# [Nick] Temporarily commenting this test out to check if it may be +# causing problems in the tests on Travis CI because of the dependence +# on ChiantiPy. This might be causing the tests on Travis CI to stall, +# while still working when we run `pytest` or `python setup.py test` +# locally. + +#@pytest.mark.parametrize('natom', natom_list) +#def test_equlibrium_state_vs_chiantipy(natom=8): +# """ +# Test equilibrium states saved in EigenData2 and compare them with +# Outputs from ChiantiPy. +# Note: +# This test requires ChiantiPy to be installed (see details +# in: https://github.com/chianti-atomic/ChiantiPy). +# """ +# try: +# import ChiantiPy.core as ch +# except ImportError: +# warnings.warn('ChiantiPy is required in this test.', UserWarning) +# return +# +# temperatures = [1.0e4, 1.0e5, 1.0e6, 1.0e7, 1.0e8] +# eqi_ch = ch.ioneq(natom) +# eqi_ch.calculate(temperatures) +# conce = eqi_ch.Ioneq +# +# table_sta = EigenData2(element=natom) +# for i in range(2): +# ch_conce = conce[:, i] +# table_conce = table_sta.equilibrium_state(T_e=temperatures[i]) +# assert ch_conce.all() == table_conce.all() +# return + +@pytest.mark.parametrize('natom', [1, 2, 6, 7, 8]) +def test_reachequlibrium_state(natom): + """ + Starting the random initial distribution, the charge states will reach + to equilibrium cases after a long time. + In this test, we set the ionization and recombination rates at + Te0=2.0e6 K and plasma density ne0=1.0e+7. A random charge states + distribution will be finally closed to equilibrium distribution at + 2.0e6K. + """ + # + # Initial conditions, set plasma temperature, density and dt + # + element_symbol = atomic.atomic_symbol(int(natom)) + te0 = 1.0e+6 + ne0 = 1.0e+8 + + # Start from any ionizaiont states, e.g., Te = 4.0d4 K, + time = 0 + table = EigenData2(element=natom) + f0 = table.equilibrium_state(T_e=4.0e4) + + print('START test_reachequlibrium_state:') + print(f'time_sta = ', time) + print(f'INI: ', f0) + print(f'Sum(f0) = ', np.sum(f0)) + + # After time + dt: + dt = 1.0e+7 + ft = func_solver_eigenval(natom, te0, ne0, time+dt, f0, table) + + print(f'time_end = ', time+dt) + print(f'NEI:', ft) + print(f'Sum(ft) = ', np.sum(ft)) + print(f'EI :', table.equilibrium_state(T_e=te0)) + print("End Test.\n") + + assert np.isclose(np.sum(ft), 1), 'np.sum(ft) is not approximately 1' + assert np.isclose(np.sum(f0), 1), 'np.sum(f0) is not approximately 1' + assert np.allclose(ft, table.equilibrium_state(T_e=te0)) + + +def test_reachequlibrium_state_multisteps(natom=8): + """ + Starting the random initial distribution, the charge states will reach + to equilibrium cases after a long time (multiple steps). + In this test, we set the ionization and recombination rates at + Te0=2.0e6 K and plasma density ne0=1.0e+7. A random charge states + distribution will be finally closed to equilibrium distribution at + 2.0e6K. + """ + # + # Initial conditions, set plasma temperature, density and dt + # + element_symbol = atomic.atomic_symbol(int(natom)) + te0 = 1.0e+6 # unit: K + ne0 = 1.0e+8 # unit: cm^-3 + + # Start from any ionizaiont states, e.g., Te = 4.0d4 K, + time = 0 + table = EigenData2(element=natom) + f0 = table.equilibrium_state(T_e=4.0e+4) + + print('START test_reachequlibrium_state_multisteps:') + print(f'time_sta = ', time) + print(f'INI: ', f0) + print(f'Sum(f0) = ', np.sum(f0)) + + # After time + dt: + dt = 100000.0 # unit: second + + # Enter the time loop: + for it in range(100): + ft = func_solver_eigenval(natom, te0, ne0, time+dt, f0, table) + f0 = np.copy(ft) + time = time+dt + + print(f'time_end = ', time+dt) + print(f'NEI:', ft) + print(f'Sum(ft) = ', np.sum(ft)) + + assert np.isclose(np.sum(ft), 1) + + print(f"EI :", table.equilibrium_state(T_e=te0)) + print("End Test.\n") + +# Temporarily test only lighter elements due to Travis CI delays + +#@pytest.mark.parametrize('atomic_numb', np.arange(1, 27)) +@pytest.mark.parametrize('atomic_numb', np.arange(1, 10)) +def test_element_range(atomic_numb): + """ + Function test_element_range: + This function is used to test element including Hydrogen to Iron. + """ + try: + element_symbol = atomic.atomic_symbol(int(atomic_numb)) + eigen = EigenData2(element=element_symbol) + except Exception as exc: + raise Exception(f"Problem with atomic number={atomic_numb}.") from exc + + eigen=0 # attempt to clear up memory From 085ea58a679eb75deeb678f6bee261f467fddc7b Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Thu, 30 Apr 2020 15:55:01 -0400 Subject: [PATCH 03/49] Add h5py dependency and update test configuration --- setup.cfg | 1 + setup.py | 2 +- tox.ini | 3 ++- 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/setup.cfg b/setup.cfg index 0d2cd85..e261197 100644 --- a/setup.cfg +++ b/setup.cfg @@ -17,6 +17,7 @@ setup_requires = install_requires = astropy plasmapy + h5py [options.entry_points] diff --git a/setup.py b/setup.py index afca96e..88616fd 100755 --- a/setup.py +++ b/setup.py @@ -20,5 +20,5 @@ "write_to": os.path.join("plasmapy_nei", "version.py"), "write_to_template": VERSION_TEMPLATE, }, - install_requires=["plasmapy>=0.3.1", "numpy>=1.16", "astropy>=3.2"], + install_requires=["plasmapy>=0.3.1", "numpy>=1.16", "astropy>=3.2", "h5py"], ) diff --git a/tox.ini b/tox.ini index 219de85..155c403 100644 --- a/tox.ini +++ b/tox.ini @@ -1,6 +1,6 @@ [tox] envlist = - py{37,38}-test + py{37}-test build_docs codestyle isolated_build = True @@ -27,6 +27,7 @@ description = deps = # The following indicates which extras_require from setup.cfg will be installed plasmapy + h5py extras = test alldeps: all From e3adaa76d3c1ceceaab9a5f664339204f2bbe56c Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Thu, 30 Apr 2020 15:58:57 -0400 Subject: [PATCH 04/49] Reformat transferred files with black --- plasmapy_nei/eigen/eigenvaluetable.py | 102 +++++++++--------- .../eigen/tests/test_eigenvaluetable.py | 91 ++++++++-------- 2 files changed, 100 insertions(+), 93 deletions(-) diff --git a/plasmapy_nei/eigen/eigenvaluetable.py b/plasmapy_nei/eigen/eigenvaluetable.py index 3cad9d5..118602f 100644 --- a/plasmapy_nei/eigen/eigenvaluetable.py +++ b/plasmapy_nei/eigen/eigenvaluetable.py @@ -52,7 +52,7 @@ class EigenData2: """ - def __init__(self, element='H'): + def __init__(self, element="H"): """Read in the """ self._element = element @@ -61,17 +61,17 @@ def __init__(self, element='H'): # # 1. Read ionization and recombination rates # - data_dir = __path__[0] + '/data/ionizrecombrates/chianti_8.07/' - filename = data_dir + 'ionrecomb_rate.h5' - f = h5py.File(filename, 'r') + data_dir = __path__[0] + "/data/ionizrecombrates/chianti_8.07/" + filename = data_dir + "ionrecomb_rate.h5" + f = h5py.File(filename, "r") atomic_numb = atomic.atomic_number(element) nstates = atomic_numb + 1 - self._temperature_grid = f['te_gird'][:] + self._temperature_grid = f["te_gird"][:] ntemp = len(self._temperature_grid) - c_ori = f['ioniz_rate'][:] - r_ori = f['recomb_rate'][:] + c_ori = f["ioniz_rate"][:] + r_ori = f["recomb_rate"][:] f.close() # @@ -80,10 +80,10 @@ def __init__(self, element='H'): c_rate = np.zeros((ntemp, nstates)) r_rate = np.zeros((ntemp, nstates)) for ite in range(ntemp): - for i in range(nstates-1): - c_rate[ite, i] = c_ori[i, atomic_numb-1, ite] + for i in range(nstates - 1): + c_rate[ite, i] = c_ori[i, atomic_numb - 1, ite] for i in range(1, nstates): - r_rate[ite, i] = r_ori[i-1, atomic_numb-1, ite] + r_rate[ite, i] = r_ori[i - 1, atomic_numb - 1, ite] # # 2. Definet the grid size @@ -95,24 +95,21 @@ def __init__(self, element='H'): # # Compute eigenvalues and eigenvectors # - self._ionization_rate = np.ndarray(shape=(ntemp, nstates), - dtype=np.float64) + self._ionization_rate = np.ndarray(shape=(ntemp, nstates), dtype=np.float64) - self._recombination_rate = np.ndarray(shape=(ntemp, nstates), - dtype=np.float64) + self._recombination_rate = np.ndarray(shape=(ntemp, nstates), dtype=np.float64) - self._equilibrium_states = np.ndarray(shape=(ntemp, nstates), - dtype=np.float64) + self._equilibrium_states = np.ndarray(shape=(ntemp, nstates), dtype=np.float64) - self._eigenvalues = np.ndarray(shape=(ntemp, nstates), - dtype=np.float64) + self._eigenvalues = np.ndarray(shape=(ntemp, nstates), dtype=np.float64) - self._eigenvectors = np.ndarray(shape=(ntemp, nstates, nstates), - dtype=np.float64) + self._eigenvectors = np.ndarray( + shape=(ntemp, nstates, nstates), dtype=np.float64 + ) self._eigenvector_inverses = np.ndarray( - shape=(ntemp, nstates, nstates), - dtype=np.float64) + shape=(ntemp, nstates, nstates), dtype=np.float64 + ) # # Save ionization and recombination rates @@ -144,18 +141,18 @@ def __init__(self, element='H'): A[ion, jon] = 0.0 # Give coefficients - for ion in range(1, nstates-1): - A[ion, ion-1] = carr[ion-1] - A[ion, ion] = -(carr[ion]+rarr[ion]) - A[ion, ion+1] = rarr[ion+1] + for ion in range(1, nstates - 1): + A[ion, ion - 1] = carr[ion - 1] + A[ion, ion] = -(carr[ion] + rarr[ion]) + A[ion, ion + 1] = rarr[ion + 1] # The first row A[0, 0] = -carr[0] A[0, 1] = rarr[1] # The last row - A[nstates-1, nstates-2] = carr[nstates-2] - A[nstates-1, nstates-1] = -rarr[nstates-1] + A[nstates - 1, nstates - 2] = carr[nstates - 2] + A[nstates - 1, nstates - 1] = -rarr[nstates - 1] # Compute eigenvalues and eigenvectors using Scipy la, v = LA.eig(A) @@ -198,14 +195,20 @@ def _get_temperature_index(self, T_e): elif T_e == T_e_grid_min: return 0 elif T_e > T_e_grid_max: - warnings.warn(f"Temperature exceeds the temperature grid " - f"boundary: temperature index will be reset " - f"to {self._ntemp - 1}", UserWarning) - return self._ntemp-1 + warnings.warn( + f"Temperature exceeds the temperature grid " + f"boundary: temperature index will be reset " + f"to {self._ntemp - 1}", + UserWarning, + ) + return self._ntemp - 1 elif T_e < T_e_grid_min: - warnings.warn(f"Temperature is below the temperature grid " - f"boundary: temperature index will be reset to " - f"0.", UserWarning) + warnings.warn( + f"Temperature is below the temperature grid " + f"boundary: temperature index will be reset to " + f"0.", + UserWarning, + ) return 0 # TODO: Add a test to check that the temperature grid is monotonic @@ -215,7 +218,7 @@ def _get_temperature_index(self, T_e): index = res_ind[0] dte_l = abs(T_e - T_e_array[index - 1]) # re-check the neighbor point dte_r = abs(T_e - T_e_array[index]) - if (dte_l <= dte_r): + if dte_l <= dte_r: index = index - 1 return index @@ -277,7 +280,6 @@ def eigenvector_inverses(self, T_e=None, T_e_index=None): else: raise AttributeError("The temperature has not been set.") - def equilibrium_state(self, T_e=None, T_e_index=None): """Returns the equilibrium charge state distribution for the temperature specified in the class.""" @@ -302,19 +304,19 @@ def _function_eqi(self, ioniz_rate, recomb_rate, natom): # The start index is 1. for i in range(nstates): - c[i+1] = ioniz_rate[i] - r[i+1] = recomb_rate[i] + c[i + 1] = ioniz_rate[i] + r[i + 1] = recomb_rate[i] # Set f1 f[1] = 1.0 # f2 = c1*f1/r2 - f[2] = c[1]*f[1]/r[2] + f[2] = c[1] * f[1] / r[2] # # For Hydrogen, the following loop is not necessary. # - if (natom <= 1): - f[1] = 1.0/(1.0 + c[1]/r[2]) - f[2] = c[1]*f[1]/r[2] + if natom <= 1: + f[1] = 1.0 / (1.0 + c[1] / r[2]) + f[2] = c[1] * f[1] / r[2] conce[0:2] = f[1:3] return conce @@ -323,20 +325,20 @@ def _function_eqi(self, ioniz_rate, recomb_rate, natom): # # f(i+1) = -(c(i-1)*f(i-1) - (c(i)+r(i)*f(i)))/r(i+1) for k in range(2, natom): - f[k+1] = (-c[k-1]*f[k-1] + (c[k]+r[k])*f[k])/r[k+1] + f[k + 1] = (-c[k - 1] * f[k - 1] + (c[k] + r[k]) * f[k]) / r[k + 1] # f(natom+1) = c(natom)*f(natom)/r(natom+1) - f[natom+1] = c[natom]*f[natom]/r[natom+1] + f[natom + 1] = c[natom] * f[natom] / r[natom + 1] # f1 = 1/sum(f(*)) - f[1] = 1.0/np.sum(f) + f[1] = 1.0 / np.sum(f) # f2 = c1*f1/r2 - f[2] = c[1]*f[1]/r[2] + f[2] = c[1] * f[1] / r[2] # f(i+1) = -(c(i-1)*f(i-1) - (c(i)+r(i)*f(i)))/r(i+1) for k in range(2, natom): - f[k+1] = (-c[k-1]*f[k-1] + (c[k]+r[k])*f[k])/r[k+1] + f[k + 1] = (-c[k - 1] * f[k - 1] + (c[k] + r[k]) * f[k]) / r[k + 1] # f(natom+1) = c(natom)*f(natom)/r(natom+1) - f[natom+1] = c[natom]*f[natom]/r[natom+1] + f[natom + 1] = c[natom] * f[natom] / r[natom + 1] # - conce[0:nstates] = f[1:nstates+1] + conce[0:nstates] = f[1 : nstates + 1] return conce diff --git a/plasmapy_nei/eigen/tests/test_eigenvaluetable.py b/plasmapy_nei/eigen/tests/test_eigenvaluetable.py index fc34396..04b85fa 100644 --- a/plasmapy_nei/eigen/tests/test_eigenvaluetable.py +++ b/plasmapy_nei/eigen/tests/test_eigenvaluetable.py @@ -5,29 +5,31 @@ from ..eigenvaluetable import EigenData2 import pytest -#------------------------------------------------------------------------------- +# ------------------------------------------------------------------------------- # Set elements and temperary list as testing inputs -#------------------------------------------------------------------------------- +# ------------------------------------------------------------------------------- natom_list = np.arange(1, 27) natom = 8 -#------------------------------------------------------------------------------ +# ------------------------------------------------------------------------------ # function: Time-Advance solover -#------------------------------------------------------------------------------ +# ------------------------------------------------------------------------------ def func_solver_eigenval(natom, te, ne, dt, f0, table): """ The testing function for performing time_advance calculations. """ common_index = table._get_temperature_index(te) - evals = table.eigenvalues(T_e_index=common_index) # find eigenvalues on the chosen Te node + evals = table.eigenvalues( + T_e_index=common_index + ) # find eigenvalues on the chosen Te node evect = table.eigenvectors(T_e_index=common_index) evect_invers = table.eigenvector_inverses(T_e_index=common_index) # define the temperary diagonal matrix - diagona_evals = np.zeros((natom+1, natom+1)) - for ii in range(0, natom+1): - diagona_evals[ii,ii] = np.exp(evals[ii]*dt*ne) + diagona_evals = np.zeros((natom + 1, natom + 1)) + for ii in range(0, natom + 1): + diagona_evals[ii, ii] = np.exp(evals[ii] * dt * ne) # matirx operation matrix_1 = np.dot(diagona_evals, evect) @@ -38,19 +40,20 @@ def func_solver_eigenval(natom, te, ne, dt, f0, table): # re-check the smallest value minconce = 1.0e-15 - for ii in np.arange(0, natom+1, dtype=np.int): - if (abs(ft[ii]) <= minconce): + for ii in np.arange(0, natom + 1, dtype=np.int): + if abs(ft[ii]) <= minconce: ft[ii] = 0.0 return ft + # [Nick] Temporarily commenting this test out to check if it may be # causing problems in the tests on Travis CI because of the dependence # on ChiantiPy. This might be causing the tests on Travis CI to stall, # while still working when we run `pytest` or `python setup.py test` # locally. -#@pytest.mark.parametrize('natom', natom_list) -#def test_equlibrium_state_vs_chiantipy(natom=8): +# @pytest.mark.parametrize('natom', natom_list) +# def test_equlibrium_state_vs_chiantipy(natom=8): # """ # Test equilibrium states saved in EigenData2 and compare them with # Outputs from ChiantiPy. @@ -76,7 +79,8 @@ def func_solver_eigenval(natom, te, ne, dt, f0, table): # assert ch_conce.all() == table_conce.all() # return -@pytest.mark.parametrize('natom', [1, 2, 6, 7, 8]) + +@pytest.mark.parametrize("natom", [1, 2, 6, 7, 8]) def test_reachequlibrium_state(natom): """ Starting the random initial distribution, the charge states will reach @@ -90,31 +94,31 @@ def test_reachequlibrium_state(natom): # Initial conditions, set plasma temperature, density and dt # element_symbol = atomic.atomic_symbol(int(natom)) - te0 = 1.0e+6 - ne0 = 1.0e+8 + te0 = 1.0e6 + ne0 = 1.0e8 # Start from any ionizaiont states, e.g., Te = 4.0d4 K, time = 0 table = EigenData2(element=natom) f0 = table.equilibrium_state(T_e=4.0e4) - print('START test_reachequlibrium_state:') - print(f'time_sta = ', time) - print(f'INI: ', f0) - print(f'Sum(f0) = ', np.sum(f0)) + print("START test_reachequlibrium_state:") + print(f"time_sta = ", time) + print(f"INI: ", f0) + print(f"Sum(f0) = ", np.sum(f0)) # After time + dt: - dt = 1.0e+7 - ft = func_solver_eigenval(natom, te0, ne0, time+dt, f0, table) + dt = 1.0e7 + ft = func_solver_eigenval(natom, te0, ne0, time + dt, f0, table) - print(f'time_end = ', time+dt) - print(f'NEI:', ft) - print(f'Sum(ft) = ', np.sum(ft)) - print(f'EI :', table.equilibrium_state(T_e=te0)) + print(f"time_end = ", time + dt) + print(f"NEI:", ft) + print(f"Sum(ft) = ", np.sum(ft)) + print(f"EI :", table.equilibrium_state(T_e=te0)) print("End Test.\n") - assert np.isclose(np.sum(ft), 1), 'np.sum(ft) is not approximately 1' - assert np.isclose(np.sum(f0), 1), 'np.sum(f0) is not approximately 1' + assert np.isclose(np.sum(ft), 1), "np.sum(ft) is not approximately 1" + assert np.isclose(np.sum(f0), 1), "np.sum(f0) is not approximately 1" assert np.allclose(ft, table.equilibrium_state(T_e=te0)) @@ -131,41 +135,42 @@ def test_reachequlibrium_state_multisteps(natom=8): # Initial conditions, set plasma temperature, density and dt # element_symbol = atomic.atomic_symbol(int(natom)) - te0 = 1.0e+6 # unit: K - ne0 = 1.0e+8 # unit: cm^-3 + te0 = 1.0e6 # unit: K + ne0 = 1.0e8 # unit: cm^-3 # Start from any ionizaiont states, e.g., Te = 4.0d4 K, time = 0 table = EigenData2(element=natom) - f0 = table.equilibrium_state(T_e=4.0e+4) + f0 = table.equilibrium_state(T_e=4.0e4) - print('START test_reachequlibrium_state_multisteps:') - print(f'time_sta = ', time) - print(f'INI: ', f0) - print(f'Sum(f0) = ', np.sum(f0)) + print("START test_reachequlibrium_state_multisteps:") + print(f"time_sta = ", time) + print(f"INI: ", f0) + print(f"Sum(f0) = ", np.sum(f0)) # After time + dt: - dt = 100000.0 # unit: second + dt = 100000.0 # unit: second # Enter the time loop: for it in range(100): - ft = func_solver_eigenval(natom, te0, ne0, time+dt, f0, table) + ft = func_solver_eigenval(natom, te0, ne0, time + dt, f0, table) f0 = np.copy(ft) - time = time+dt + time = time + dt - print(f'time_end = ', time+dt) - print(f'NEI:', ft) - print(f'Sum(ft) = ', np.sum(ft)) + print(f"time_end = ", time + dt) + print(f"NEI:", ft) + print(f"Sum(ft) = ", np.sum(ft)) assert np.isclose(np.sum(ft), 1) print(f"EI :", table.equilibrium_state(T_e=te0)) print("End Test.\n") + # Temporarily test only lighter elements due to Travis CI delays -#@pytest.mark.parametrize('atomic_numb', np.arange(1, 27)) -@pytest.mark.parametrize('atomic_numb', np.arange(1, 10)) +# @pytest.mark.parametrize('atomic_numb', np.arange(1, 27)) +@pytest.mark.parametrize("atomic_numb", np.arange(1, 10)) def test_element_range(atomic_numb): """ Function test_element_range: @@ -177,4 +182,4 @@ def test_element_range(atomic_numb): except Exception as exc: raise Exception(f"Problem with atomic number={atomic_numb}.") from exc - eigen=0 # attempt to clear up memory + eigen = 0 # attempt to clear up memory From 2dd31fae0c0352a066d0959086bf5fb57383a865 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Thu, 30 Apr 2020 16:03:41 -0400 Subject: [PATCH 05/49] Skip doctests and mark that copied over tests will fail These tests will not work since they are in files that were directly copied over from the NEI-modeling/NEI repository. --- plasmapy_nei/eigen/eigenvaluetable.py | 12 ++++++------ plasmapy_nei/eigen/tests/test_eigenvaluetable.py | 4 +++- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/plasmapy_nei/eigen/eigenvaluetable.py b/plasmapy_nei/eigen/eigenvaluetable.py index 118602f..241b22f 100644 --- a/plasmapy_nei/eigen/eigenvaluetable.py +++ b/plasmapy_nei/eigen/eigenvaluetable.py @@ -31,23 +31,23 @@ class EigenData2: ---------- To get the table for element 'Helium' at Te=5.0e5K: - >>> table = EigenData2(element=2) - >>> table.temperature=5.0e5 + >>> table = EigenData2(element=2) # doctest: +SKIP + >>> table.temperature=5.0e5 # doctest: +SKIP Output eigenvals: - >>> table.eigenvalues() + >>> table.eigenvalues() # doctest: +SKIP array([-1.12343827e-08, -9.00005970e-10, 8.58945565e-30]) Output equilibrium states: - >>> table.equilibrium_state() + >>> table.equilibrium_state() # doctest: +SKIP array([9.57162006e-09, 1.26564673e-04, 9.99873426e-01]) Or one may output properties at other temperatures, for example: - >>> table.eigenvalues(T_e=12589.254117941662) + >>> table.eigenvalues(T_e=12589.254117941662) # doctest: +SKIP array([-1.87480784e-12, -3.73191378e-13, 0.00000000e+00]) Or input temperature index on the grid: - >>> table.eigenvalues(T_e_index=100) + >>> table.eigenvalues(T_e_index=100) # doctest: +SKIP array([-1.87480784e-12, -3.73191378e-13, 0.00000000e+00]) """ diff --git a/plasmapy_nei/eigen/tests/test_eigenvaluetable.py b/plasmapy_nei/eigen/tests/test_eigenvaluetable.py index 04b85fa..83ac689 100644 --- a/plasmapy_nei/eigen/tests/test_eigenvaluetable.py +++ b/plasmapy_nei/eigen/tests/test_eigenvaluetable.py @@ -79,7 +79,7 @@ def func_solver_eigenval(natom, te, ne, dt, f0, table): # assert ch_conce.all() == table_conce.all() # return - +@pytest.mark.xfail @pytest.mark.parametrize("natom", [1, 2, 6, 7, 8]) def test_reachequlibrium_state(natom): """ @@ -122,6 +122,7 @@ def test_reachequlibrium_state(natom): assert np.allclose(ft, table.equilibrium_state(T_e=te0)) +@pytest.mark.xfail def test_reachequlibrium_state_multisteps(natom=8): """ Starting the random initial distribution, the charge states will reach @@ -170,6 +171,7 @@ def test_reachequlibrium_state_multisteps(natom=8): # Temporarily test only lighter elements due to Travis CI delays # @pytest.mark.parametrize('atomic_numb', np.arange(1, 27)) +@pytest.mark.xfail @pytest.mark.parametrize("atomic_numb", np.arange(1, 10)) def test_element_range(atomic_numb): """ From 83c52269d6657aaae1e0c4fbe3b42834030a57c2 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Thu, 30 Apr 2020 17:13:46 -0400 Subject: [PATCH 06/49] Set include_package_data to True in setup.py --- setup.py | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.py b/setup.py index 88616fd..6fdd635 100755 --- a/setup.py +++ b/setup.py @@ -21,4 +21,5 @@ "write_to_template": VERSION_TEMPLATE, }, install_requires=["plasmapy>=0.3.1", "numpy>=1.16", "astropy>=3.2", "h5py"], + include_package_data=True, ) From bb18e39fd0c5f879c0be977786d7b05f125c35cb Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Thu, 30 Apr 2020 20:26:59 -0400 Subject: [PATCH 07/49] Add EigenData class and tests The class is adapted from the class EigenData2 which is in the eigenvaluetable.py file that will be deleted soon. The tests are adapted from the test_eigenvaluetable.py class that will also be deleted soon. --- plasmapy_nei/eigen/__init__.py | 2 + plasmapy_nei/eigen/eigenclass.py | 354 ++++++++++++++++++++ plasmapy_nei/eigen/tests/test_eigenclass.py | 135 ++++++++ 3 files changed, 491 insertions(+) create mode 100644 plasmapy_nei/eigen/eigenclass.py create mode 100644 plasmapy_nei/eigen/tests/test_eigenclass.py diff --git a/plasmapy_nei/eigen/__init__.py b/plasmapy_nei/eigen/__init__.py index 77f744f..6844556 100644 --- a/plasmapy_nei/eigen/__init__.py +++ b/plasmapy_nei/eigen/__init__.py @@ -1 +1,3 @@ """Classes for accessing eigentables for ionization and recombination rates.""" + +from .eigenclass import EigenData diff --git a/plasmapy_nei/eigen/eigenclass.py b/plasmapy_nei/eigen/eigenclass.py new file mode 100644 index 0000000..43b4919 --- /dev/null +++ b/plasmapy_nei/eigen/eigenclass.py @@ -0,0 +1,354 @@ +""" +Contains the class to provide access to eigenvalue data for the +ionization and recombination rates. +""" + +__all__ = ["EigenData"] + +import h5py +import numpy as np +from numpy import linalg as LA +import pkg_resources +import warnings + +try: + from plasmapy.particles import Particle, particle_input +except (ImportError, ModuleNotFoundError): + from plasmapy.atomic import Particle, particle_input + +max_atomic_number = 30 # TODO: double check if this is the correct number + + +def _get_equilibrium_charge_states(ioniz_rate, recomb_rate, natom): + """ + Compute the equilibrium charge state distribution for the + temperature specified using ionization and recombination rates. + + Parameters + ---------- + """ + nstates = natom + 1 + conce = np.zeros(nstates) + f = np.zeros(nstates + 1) + c = np.zeros(nstates + 1) + r = np.zeros(nstates + 1) + + # The start index is 1. + for i in range(nstates): + c[i + 1] = ioniz_rate[i] + r[i + 1] = recomb_rate[i] + + f[1] = 1.0 + f[2] = c[1] * f[1] / r[2] + + # simplest for hydrogen + if natom == 1: + f[1] = 1.0 / (1.0 + c[1] / r[2]) + f[2] = c[1] * f[1] / r[2] + conce[0:2] = f[1:3] + return conce + + # for other elements + + for k in range(2, natom): + f[k + 1] = (-c[k - 1] * f[k - 1] + (c[k] + r[k]) * f[k]) / r[k + 1] + + f[natom + 1] = c[natom] * f[natom] / r[natom + 1] + + f[1] = 1.0 / np.sum(f) + + f[2] = c[1] * f[1] / r[2] + + for k in range(2, natom): + f[k + 1] = (-c[k - 1] * f[k - 1] + (c[k] + r[k]) * f[k]) / r[k + 1] + + f[natom + 1] = c[natom] * f[natom] / r[natom + 1] + + conce[0:nstates] = f[1 : nstates + 1] + return conce + + +class EigenData: + """ + Provides access to ionization and recombination rate data. + + Parameters + ---------- + element : particle-like + Representation of the element to access data for. + + Examples + -------- + >>> eigendata = EigenData("He") + >>> eigendata.nstates + 3 + """ + + def _validate_element(self, element): + + # The following might not be needed if the check is in @particle_input + if not element.is_category(require="element", exclude=["ion", "isotope"]): + raise ValueError(f"{element} is not an element") + + if element.atomic_number > max_atomic_number: + raise ValueError("Need an element") + + self._element = element + + def _load_data(self): + """ + Retrieve data from the HDF5 file containing ionization and + recombination rates. + """ + filename = pkg_resources.resource_filename( + "plasmapy_nei", + "data/ionrecomb_rate.h5", # from Chianti database, version 8.07 + ) + + # TODO: enable different HDF5 files + + with h5py.File(filename, "r") as file: + self._temperature_grid = file["te_gird"][:] # TODO: fix typo in HDF5 file + self._ntemp = len(self.temperature_grid) + c_ori = file["ioniz_rate"][:] + r_ori = file["recomb_rate"][:] + + c_rate = np.zeros((self.ntemp, self.nstates)) + r_rate = np.zeros((self.ntemp, self.nstates)) + for temperature_index in range(self.ntemp): + for i in range(self.nstates - 1): + c_rate[temperature_index, i] = c_ori[ + i, self.element.atomic_number - 1, temperature_index + ] + for i in range(1, self.nstates): + r_rate[temperature_index, i] = r_ori[ + i - 1, self.element.atomic_number - 1, temperature_index + ] + + self._ionization_rate = np.ndarray( + shape=(self.ntemp, self.nstates), dtype=np.float64 + ) + self._recombination_rate = np.ndarray( + shape=(self.ntemp, self.nstates), dtype=np.float64 + ) + self._equilibrium_states = np.ndarray( + shape=(self.ntemp, self.nstates), dtype=np.float64 + ) + self._eigenvalues = np.ndarray( + shape=(self.ntemp, self.nstates), dtype=np.float64 + ) + + self._eigenvectors = np.ndarray( + shape=(self.ntemp, self.nstates, self.nstates), dtype=np.float64, + ) + + self._eigenvector_inverses = np.ndarray( + shape=(self.ntemp, self.nstates, self.nstates), dtype=np.float64, + ) + + self._ionization_rate[:, :] = c_rate[:, :] + self._recombination_rate[:, :] = r_rate[:, :] + + # Define the coefficients matrix A. The first dimension is + # for elements, and the second number of equations. + + nequations = self.nstates + coefficients_matrix = np.zeros( + shape=(self.nstates, nequations), dtype=np.float64 + ) + + # Enter temperature loop over the whole temperature grid + + for temperature_index in range(self.ntemp): + + # Ionization and recombination rate at Te(temperature_index) + carr = c_rate[temperature_index, :] + rarr = r_rate[temperature_index, :] + + eqi = _get_equilibrium_charge_states( + carr, rarr, self.element.atomic_number, + ) + + for ion in range(1, self.nstates - 1): + coefficients_matrix[ion, ion - 1] = carr[ion - 1] + coefficients_matrix[ion, ion] = -(carr[ion] + rarr[ion]) + coefficients_matrix[ion, ion + 1] = rarr[ion + 1] + + # The first row + coefficients_matrix[0, 0] = -carr[0] + coefficients_matrix[0, 1] = rarr[1] + + # The last row + coefficients_matrix[self.nstates - 1, self.nstates - 2] = carr[ + self.nstates - 2 + ] + coefficients_matrix[self.nstates - 1, self.nstates - 1] = -rarr[ + self.nstates - 1 + ] + + # Compute eigenvalues and eigenvectors + eigenvalues, eigenvectors = LA.eig(coefficients_matrix) + + # Rearrange the eigenvalues. + idx = np.argsort(eigenvalues) + eigenvalues = eigenvalues[idx] + eigenvectors = eigenvectors[:, idx] + + # Compute inverse of eigenvectors + v_inverse = LA.inv(eigenvectors) + + # transpose the order to as same as the Fortran Version + eigenvectors = eigenvectors.transpose() + v_inverse = v_inverse.transpose() + + # Save eigenvalues and eigenvectors into arrays + for j in range(self.nstates): + self._eigenvalues[temperature_index, j] = eigenvalues[j] + self._equilibrium_states[temperature_index, j] = eqi[j] + for i in range(self.nstates): + self._eigenvectors[temperature_index, i, j] = eigenvectors[i, j] + self._eigenvector_inverses[temperature_index, i, j] = v_inverse[ + i, j + ] + + @particle_input + def __init__(self, element: Particle): + + self._validate_element(element) + self._load_data() + + def _get_temperature_index(self, T_e): # TODO: extract this to a function + """Return the temperature index closest to a particular temperature.""" + T_e_array = self._temperature_grid + + # Check the temperature range + T_e_grid_max = np.amax(T_e_array) + T_e_grid_min = np.amin(T_e_array) + + if T_e == T_e_grid_max: + return self._ntemp - 1 + elif T_e == T_e_grid_min: + return 0 + elif T_e > T_e_grid_max: + warnings.warn( + f"Temperature exceeds the temperature grid " + f"boundary: temperature index will be reset " + f"to {self._ntemp - 1}", + UserWarning, + ) + return self._ntemp - 1 + elif T_e < T_e_grid_min: + warnings.warn( + f"Temperature is below the temperature grid " + f"boundary: temperature index will be reset to " + f"0.", + UserWarning, + ) + return 0 + + # TODO: Add a test to check that the temperature grid is monotonic + + res = np.where(T_e_array >= T_e) + res_ind = res[0] + index = res_ind[0] + dte_l = abs(T_e - T_e_array[index - 1]) # re-check the neighbor point + dte_r = abs(T_e - T_e_array[index]) + if dte_l <= dte_r: + index = index - 1 + return index + + @property + def temperature(self): + return self._temperature + + @property + def temperature(self, T_e): + self._temperature = T_e + self._te_index = self._get_temperature_index(T_e) + + @property + def temperature_grid(self): + return self._temperature_grid + + @property + def element(self) -> Particle: + """The element corresponding to an instance of this class.""" + return self._element + + @property + def nstates(self) -> int: + """Number of charge states corresponding to the element.""" + return self.element.atomic_number + 1 + + @property + def ntemp(self) -> int: + """Number of points in ``temperature_grid``.""" + return self._ntemp + + @property + def ionization_rate(self): + return self._ionization_rate + + @property + def recombination_rate(self): + return self._recombination_rate + + def eigenvalues(self, T_e=None, T_e_index=None): + """ + Returns the eigenvalues for the ionization and recombination + rates for the temperature specified in the class. + """ + if T_e_index: + return self._eigenvalues[T_e_index, :] + elif T_e: + T_e_index = self._get_temperature_index(T_e) + return self._eigenvalues[T_e_index, :] + elif self.temperature: + return self._eigenvalues[self._te_index, :] + else: + raise AttributeError("The temperature has not been set.") + + def eigenvectors(self, T_e=None, T_e_index=None): + """ + Returns the eigenvectors for the ionization and recombination + rates for the temperature specified in the class. + """ + if T_e_index: + return self._eigenvectors[T_e_index, :, :] + elif T_e: + T_e_index = self._get_temperature_index(T_e) + return self._eigenvectors[T_e_index, :, :] + elif self.temperature: + return self._eigenvectors[self._te_index, :, :] + else: + raise AttributeError("The temperature has not been set.") + + def eigenvector_inverses(self, T_e=None, T_e_index=None): + """ + Returns the inverses of the eigenvectors for the ionization and + recombination rates for the temperature specified in the class. + """ + if T_e_index: + return self._eigenvector_inverses[T_e_index, :, :] + elif T_e: + T_e_index = self._get_temperature_index(T_e) + return self._eigenvector_inverses[T_e_index, :, :] + elif self.temperature: + return self._eigenvector_inverses[self._te_index, :, :] + else: + raise AttributeError("The temperature has not been set.") + + def equilibrium_state(self, T_e=None, T_e_index=None): + """ + Return the equilibrium charge state distribution for the + temperature specified in the class. + """ + if T_e_index: + return self._equilibrium_states[T_e_index, :] + elif T_e: + T_e_index = self._get_temperature_index(T_e) + return self._equilibrium_states[T_e_index, :] + elif self.temperature: + return self._equilibrium_states[self._te_index, :] + else: + raise AttributeError("The temperature has not been set.") diff --git a/plasmapy_nei/eigen/tests/test_eigenclass.py b/plasmapy_nei/eigen/tests/test_eigenclass.py new file mode 100644 index 0000000..18db8dd --- /dev/null +++ b/plasmapy_nei/eigen/tests/test_eigenclass.py @@ -0,0 +1,135 @@ +"""Tests of the class to store package data.""" + +from plasmapy import atomic +from plasmapy_nei.eigen import EigenData +import pytest +import numpy as np + + +@pytest.mark.parametrize("atomic_numb", np.arange(1, 30)) +def test_instantiation(atomic_numb): + try: + element_symbol = atomic.atomic_symbol(int(atomic_numb)) + EigenData(element=element_symbol) + except Exception as exc: + raise Exception( + f"Problem with instantiation for atomic number={atomic_numb}." + ) from exc + + +def time_advance_solver_for_testing(natom, te, ne, dt, f0, table): + """Testing function for performing time advance calculations""" + + common_index = table._get_temperature_index(te) + evals = table.eigenvalues( + T_e_index=common_index + ) # find eigenvalues on the chosen Te node + evect = table.eigenvectors(T_e_index=common_index) + evect_invers = table.eigenvector_inverses(T_e_index=common_index) + + # define the temperary diagonal matrix + diagona_evals = np.zeros((natom + 1, natom + 1)) + for ii in range(0, natom + 1): + diagona_evals[ii, ii] = np.exp(evals[ii] * dt * ne) + + # matrix operation + matrix_1 = np.dot(diagona_evals, evect) + matrix_2 = np.dot(evect_invers, matrix_1) + + # get ionic fraction at (time+dt) + ft = np.dot(f0, matrix_2) + + # re-check the smallest value + minconce = 1.0e-15 + for ii in np.arange(0, natom + 1, dtype=np.int): + if abs(ft[ii]) <= minconce: + ft[ii] = 0.0 + return ft + + +@pytest.mark.parametrize("natom", [1, 2, 6, 7, 8]) +def test_reachequlibrium_state(natom): + """ + Starting the random initial distribution, the charge states will reach + to equilibrium cases after a long time. + In this test, we set the ionization and recombination rates at + Te0=2.0e6 K and plasma density ne0=1.0e+7. A random charge states + distribution will be finally closed to equilibrium distribution at + 2.0e6K. + """ + # + # Initial conditions, set plasma temperature, density and dt + # + element_symbol = atomic.atomic_symbol(int(natom)) + te0 = 1.0e6 + ne0 = 1.0e8 + + # Start from any ionizaiont states, e.g., Te = 4.0d4 K, + time = 0 + table = EigenData(element=natom) + f0 = table.equilibrium_state(T_e=4.0e4) + + print("START test_reachequlibrium_state:") + print(f"time_sta = ", time) + print(f"INI: ", f0) + print(f"Sum(f0) = ", np.sum(f0)) + + # After time + dt: + dt = 1.0e7 + ft = time_advance_solver_for_testing(natom, te0, ne0, time + dt, f0, table) + + print(f"time_end = ", time + dt) + print(f"NEI:", ft) + print(f"Sum(ft) = ", np.sum(ft)) + print(f"EI :", table.equilibrium_state(T_e=te0)) + print("End Test.\n") + + assert np.isclose(np.sum(ft), 1), "np.sum(ft) is not approximately 1" + assert np.isclose(np.sum(f0), 1), "np.sum(f0) is not approximately 1" + assert np.allclose(ft, table.equilibrium_state(T_e=te0)) + + +def test_reachequlibrium_state_multisteps(natom=8): + """ + Starting the random initial distribution, the charge states will reach + to equilibrium cases after a long time (multiple steps). + In this test, we set the ionization and recombination rates at + Te0=2.0e6 K and plasma density ne0=1.0e+7. A random charge states + distribution will be finally closed to equilibrium distribution at + 2.0e6K. + """ + # + # Initial conditions, set plasma temperature, density and dt + # + te0 = 1.0e6 # unit: K + ne0 = 1.0e8 # unit: cm^-3 + + # Start from any ionization state, e.g., Te = 4.0d4 K, + time = 0 + table = EigenData(element=natom) + f0 = table.equilibrium_state(T_e=4.0e4) + + # print(f"time_sta = ", time) + # print(f"INI: ", f0) + # print(f"Sum(f0) = ", np.sum(f0)) + + # After time + dt: + dt = 100000.0 # unit: second + + # Enter the time loop: + for it in range(100): + ft = time_advance_solver_for_testing(natom, te0, ne0, time + dt, f0, table) + f0 = np.copy(ft) + time = time + dt + + # print(f"time_end = ", time + dt) + # print(f"NEI:", ft) + # print(f"Sum(ft) = ", np.sum(ft)) + + assert np.isclose(np.sum(ft), 1) + + # print(f"EI :", table.equilibrium_state(T_e=te0)) + # print("End Test.\n") + + +# TODO: Test that appropriate exceptions are raised for invalid inputs From d94ad98eb928e4cb23f2d049f562e940afe1872d Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Thu, 30 Apr 2020 20:29:42 -0400 Subject: [PATCH 08/49] Delete old EigenData2 class and tests These were copied over from the NEI-modeling/NEI repository. --- plasmapy_nei/eigen/eigenvaluetable.py | 344 ------------------ .../eigen/tests/test_eigenvaluetable.py | 187 ---------- 2 files changed, 531 deletions(-) delete mode 100644 plasmapy_nei/eigen/eigenvaluetable.py delete mode 100644 plasmapy_nei/eigen/tests/test_eigenvaluetable.py diff --git a/plasmapy_nei/eigen/eigenvaluetable.py b/plasmapy_nei/eigen/eigenvaluetable.py deleted file mode 100644 index 241b22f..0000000 --- a/plasmapy_nei/eigen/eigenvaluetable.py +++ /dev/null @@ -1,344 +0,0 @@ -"""The EigenData2 class.""" - -import warnings -import numpy as np -from numpy import linalg as LA -from plasmapy import atomic -from .. import __path__ -import h5py - - -class EigenData2: - """ - - A class to obtain eigenvalue tables used in NEI calculations. - - Parameters - ---------- - element : `str` or `int` - A string representing a element symbol. The defaut value - is for Hydrogen. - - Raises: - ---------- - - Properties: - temperature: `float` - The temperary is set to get the relative rates and eigen values. - In unit of K. - - Examples: - ---------- - To get the table for element 'Helium' at Te=5.0e5K: - - >>> table = EigenData2(element=2) # doctest: +SKIP - >>> table.temperature=5.0e5 # doctest: +SKIP - - Output eigenvals: - >>> table.eigenvalues() # doctest: +SKIP - array([-1.12343827e-08, -9.00005970e-10, 8.58945565e-30]) - - Output equilibrium states: - >>> table.equilibrium_state() # doctest: +SKIP - array([9.57162006e-09, 1.26564673e-04, 9.99873426e-01]) - - Or one may output properties at other temperatures, for example: - >>> table.eigenvalues(T_e=12589.254117941662) # doctest: +SKIP - array([-1.87480784e-12, -3.73191378e-13, 0.00000000e+00]) - - Or input temperature index on the grid: - >>> table.eigenvalues(T_e_index=100) # doctest: +SKIP - array([-1.87480784e-12, -3.73191378e-13, 0.00000000e+00]) - - """ - - def __init__(self, element="H"): - """Read in the """ - - self._element = element - self._temperature = None - - # - # 1. Read ionization and recombination rates - # - data_dir = __path__[0] + "/data/ionizrecombrates/chianti_8.07/" - filename = data_dir + "ionrecomb_rate.h5" - f = h5py.File(filename, "r") - - atomic_numb = atomic.atomic_number(element) - nstates = atomic_numb + 1 - - self._temperature_grid = f["te_gird"][:] - ntemp = len(self._temperature_grid) - c_ori = f["ioniz_rate"][:] - r_ori = f["recomb_rate"][:] - f.close() - - # - # Ionization and recombination rate for the current element - # - c_rate = np.zeros((ntemp, nstates)) - r_rate = np.zeros((ntemp, nstates)) - for ite in range(ntemp): - for i in range(nstates - 1): - c_rate[ite, i] = c_ori[i, atomic_numb - 1, ite] - for i in range(1, nstates): - r_rate[ite, i] = r_ori[i - 1, atomic_numb - 1, ite] - - # - # 2. Definet the grid size - # - self._ntemp = ntemp - self._atomic_numb = atomic_numb - self._nstates = nstates - - # - # Compute eigenvalues and eigenvectors - # - self._ionization_rate = np.ndarray(shape=(ntemp, nstates), dtype=np.float64) - - self._recombination_rate = np.ndarray(shape=(ntemp, nstates), dtype=np.float64) - - self._equilibrium_states = np.ndarray(shape=(ntemp, nstates), dtype=np.float64) - - self._eigenvalues = np.ndarray(shape=(ntemp, nstates), dtype=np.float64) - - self._eigenvectors = np.ndarray( - shape=(ntemp, nstates, nstates), dtype=np.float64 - ) - - self._eigenvector_inverses = np.ndarray( - shape=(ntemp, nstates, nstates), dtype=np.float64 - ) - - # - # Save ionization and recombination rates - # - self._ionization_rate = c_rate - self._recombination_rate = r_rate - - # - # Define the coefficients matrix A. The first dimension is - # for elements, and the second number of equations. - # - neqs = nstates - A = np.ndarray(shape=(nstates, neqs), dtype=np.float64) - - # - # Enter temperature loop over the whole temperature grid - # - for ite in range(ntemp): - # Ionization and recombination rate at Te(ite) - carr = c_rate[ite, :] - rarr = r_rate[ite, :] - - # Equilibirum - eqi = self._function_eqi(carr, rarr, atomic_numb) - - # Initialize A to zero - for ion in range(nstates): - for jon in range(nstates): - A[ion, jon] = 0.0 - - # Give coefficients - for ion in range(1, nstates - 1): - A[ion, ion - 1] = carr[ion - 1] - A[ion, ion] = -(carr[ion] + rarr[ion]) - A[ion, ion + 1] = rarr[ion + 1] - - # The first row - A[0, 0] = -carr[0] - A[0, 1] = rarr[1] - - # The last row - A[nstates - 1, nstates - 2] = carr[nstates - 2] - A[nstates - 1, nstates - 1] = -rarr[nstates - 1] - - # Compute eigenvalues and eigenvectors using Scipy - la, v = LA.eig(A) - - # Rerange the eigenvalues. Try a simple way in here. - idx = np.argsort(la) - la = la[idx] - v = v[:, idx] - - # Compute inverse of eigenvectors - v_inverse = LA.inv(v) - - # transpose the order to as same as the Fortran Version - v = v.transpose() - v_inverse = v_inverse.transpose() - - # Save eigenvalues and eigenvectors into arrays - for j in range(nstates): - self._eigenvalues[ite, j] = la[j] - self._equilibrium_states[ite, j] = eqi[j] - for i in range(nstates): - self._eigenvectors[ite, i, j] = v[i, j] - self._eigenvector_inverses[ite, i, j] = v_inverse[i, j] - - # - # The following Functions is used to obtain the eigen values and relative - # def properties. - # - def _get_temperature_index(self, T_e): - """Returns the temperature index closest to a particular - temperature.""" - T_e_array = self._temperature_grid - - # Check the temperature range - T_e_grid_max = np.amax(T_e_array) - T_e_grid_min = np.amin(T_e_array) - - if T_e == T_e_grid_max: - return self._ntemp - 1 - elif T_e == T_e_grid_min: - return 0 - elif T_e > T_e_grid_max: - warnings.warn( - f"Temperature exceeds the temperature grid " - f"boundary: temperature index will be reset " - f"to {self._ntemp - 1}", - UserWarning, - ) - return self._ntemp - 1 - elif T_e < T_e_grid_min: - warnings.warn( - f"Temperature is below the temperature grid " - f"boundary: temperature index will be reset to " - f"0.", - UserWarning, - ) - return 0 - - # TODO: Add a test to check that the temperature grid is monotonic - - res = np.where(T_e_array >= T_e) - res_ind = res[0] - index = res_ind[0] - dte_l = abs(T_e - T_e_array[index - 1]) # re-check the neighbor point - dte_r = abs(T_e - T_e_array[index]) - if dte_l <= dte_r: - index = index - 1 - return index - - @property - def temperature(self): - """Returns the electron temperature currently in use by this class, - or None if the temperature has not been set.""" - return self._temperature - - @temperature.setter - def temperature(self, T_e): - """Sets the electron temperature and index on the temperature grid - to be used by this class""" - # TODO: Add checks for the temperature - self._temperature = T_e - self._te_index = self._get_temperature_index(T_e) - - @property - def temperature_grid(self): - """Returns the grid of temperatures corresponding to the eigendata.""" - return self._temperature_grid - - def eigenvalues(self, T_e=None, T_e_index=None): - """Returns the eigenvalues for the ionization and recombination - rates for the temperature specified in the class.""" - if T_e_index: - return self._eigenvalues[T_e_index, :] - elif T_e: - T_e_index = self._get_temperature_index(T_e) - return self._eigenvalues[T_e_index, :] - elif self.temperature: - return self._eigenvalues[self._te_index, :] - else: - raise AttributeError("The temperature has not been set.") - - def eigenvectors(self, T_e=None, T_e_index=None): - """Returns the eigenvectors for the ionization and recombination - rates for the temperature specified in the class.""" - if T_e_index: - return self._eigenvectors[T_e_index, :, :] - elif T_e: - T_e_index = self._get_temperature_index(T_e) - return self._eigenvectors[T_e_index, :, :] - elif self.temperature: - return self._eigenvectors[self._te_index, :, :] - else: - raise AttributeError("The temperature has not been set.") - - def eigenvector_inverses(self, T_e=None, T_e_index=None): - """Returns the inverses of the eigenvectors for the ionization and - recombination rates for the temperature specified in the class.""" - if T_e_index: - return self._eigenvector_inverses[T_e_index, :, :] - elif T_e: - T_e_index = self._get_temperature_index(T_e) - return self._eigenvector_inverses[T_e_index, :, :] - elif self.temperature: - return self._eigenvector_inverses[self._te_index, :, :] - else: - raise AttributeError("The temperature has not been set.") - - def equilibrium_state(self, T_e=None, T_e_index=None): - """Returns the equilibrium charge state distribution for the - temperature specified in the class.""" - if T_e_index: - return self._equilibrium_states[T_e_index, :] - elif T_e: - T_e_index = self._get_temperature_index(T_e) - return self._equilibrium_states[T_e_index, :] - elif self.temperature: - return self._equilibrium_states[self._te_index, :] - else: - raise AttributeError("The temperature has not been set.") - - def _function_eqi(self, ioniz_rate, recomb_rate, natom): - """Compute the equilibrium charge state distribution for the - temperature specified using ionization and recombinaiton rates.""" - nstates = natom + 1 - conce = np.zeros(nstates) - f = np.zeros(nstates + 1) - c = np.zeros(nstates + 1) - r = np.zeros(nstates + 1) - - # The start index is 1. - for i in range(nstates): - c[i + 1] = ioniz_rate[i] - r[i + 1] = recomb_rate[i] - - # Set f1 - f[1] = 1.0 - # f2 = c1*f1/r2 - f[2] = c[1] * f[1] / r[2] - # - # For Hydrogen, the following loop is not necessary. - # - if natom <= 1: - f[1] = 1.0 / (1.0 + c[1] / r[2]) - f[2] = c[1] * f[1] / r[2] - conce[0:2] = f[1:3] - return conce - - # - # For other elements with atomic number >= 2: - # - # f(i+1) = -(c(i-1)*f(i-1) - (c(i)+r(i)*f(i)))/r(i+1) - for k in range(2, natom): - f[k + 1] = (-c[k - 1] * f[k - 1] + (c[k] + r[k]) * f[k]) / r[k + 1] - - # f(natom+1) = c(natom)*f(natom)/r(natom+1) - f[natom + 1] = c[natom] * f[natom] / r[natom + 1] - # f1 = 1/sum(f(*)) - f[1] = 1.0 / np.sum(f) - # f2 = c1*f1/r2 - f[2] = c[1] * f[1] / r[2] - # f(i+1) = -(c(i-1)*f(i-1) - (c(i)+r(i)*f(i)))/r(i+1) - for k in range(2, natom): - f[k + 1] = (-c[k - 1] * f[k - 1] + (c[k] + r[k]) * f[k]) / r[k + 1] - - # f(natom+1) = c(natom)*f(natom)/r(natom+1) - f[natom + 1] = c[natom] * f[natom] / r[natom + 1] - # - conce[0:nstates] = f[1 : nstates + 1] - return conce diff --git a/plasmapy_nei/eigen/tests/test_eigenvaluetable.py b/plasmapy_nei/eigen/tests/test_eigenvaluetable.py deleted file mode 100644 index 83ac689..0000000 --- a/plasmapy_nei/eigen/tests/test_eigenvaluetable.py +++ /dev/null @@ -1,187 +0,0 @@ -"""Test_eigenvaluetable""" -import warnings -import numpy as np -from plasmapy import atomic -from ..eigenvaluetable import EigenData2 -import pytest - -# ------------------------------------------------------------------------------- -# Set elements and temperary list as testing inputs -# ------------------------------------------------------------------------------- -natom_list = np.arange(1, 27) -natom = 8 - -# ------------------------------------------------------------------------------ -# function: Time-Advance solover -# ------------------------------------------------------------------------------ -def func_solver_eigenval(natom, te, ne, dt, f0, table): - """ - The testing function for performing time_advance calculations. - """ - - common_index = table._get_temperature_index(te) - evals = table.eigenvalues( - T_e_index=common_index - ) # find eigenvalues on the chosen Te node - evect = table.eigenvectors(T_e_index=common_index) - evect_invers = table.eigenvector_inverses(T_e_index=common_index) - - # define the temperary diagonal matrix - diagona_evals = np.zeros((natom + 1, natom + 1)) - for ii in range(0, natom + 1): - diagona_evals[ii, ii] = np.exp(evals[ii] * dt * ne) - - # matirx operation - matrix_1 = np.dot(diagona_evals, evect) - matrix_2 = np.dot(evect_invers, matrix_1) - - # get ions fraction at (time+dt) - ft = np.dot(f0, matrix_2) - - # re-check the smallest value - minconce = 1.0e-15 - for ii in np.arange(0, natom + 1, dtype=np.int): - if abs(ft[ii]) <= minconce: - ft[ii] = 0.0 - return ft - - -# [Nick] Temporarily commenting this test out to check if it may be -# causing problems in the tests on Travis CI because of the dependence -# on ChiantiPy. This might be causing the tests on Travis CI to stall, -# while still working when we run `pytest` or `python setup.py test` -# locally. - -# @pytest.mark.parametrize('natom', natom_list) -# def test_equlibrium_state_vs_chiantipy(natom=8): -# """ -# Test equilibrium states saved in EigenData2 and compare them with -# Outputs from ChiantiPy. -# Note: -# This test requires ChiantiPy to be installed (see details -# in: https://github.com/chianti-atomic/ChiantiPy). -# """ -# try: -# import ChiantiPy.core as ch -# except ImportError: -# warnings.warn('ChiantiPy is required in this test.', UserWarning) -# return -# -# temperatures = [1.0e4, 1.0e5, 1.0e6, 1.0e7, 1.0e8] -# eqi_ch = ch.ioneq(natom) -# eqi_ch.calculate(temperatures) -# conce = eqi_ch.Ioneq -# -# table_sta = EigenData2(element=natom) -# for i in range(2): -# ch_conce = conce[:, i] -# table_conce = table_sta.equilibrium_state(T_e=temperatures[i]) -# assert ch_conce.all() == table_conce.all() -# return - -@pytest.mark.xfail -@pytest.mark.parametrize("natom", [1, 2, 6, 7, 8]) -def test_reachequlibrium_state(natom): - """ - Starting the random initial distribution, the charge states will reach - to equilibrium cases after a long time. - In this test, we set the ionization and recombination rates at - Te0=2.0e6 K and plasma density ne0=1.0e+7. A random charge states - distribution will be finally closed to equilibrium distribution at - 2.0e6K. - """ - # - # Initial conditions, set plasma temperature, density and dt - # - element_symbol = atomic.atomic_symbol(int(natom)) - te0 = 1.0e6 - ne0 = 1.0e8 - - # Start from any ionizaiont states, e.g., Te = 4.0d4 K, - time = 0 - table = EigenData2(element=natom) - f0 = table.equilibrium_state(T_e=4.0e4) - - print("START test_reachequlibrium_state:") - print(f"time_sta = ", time) - print(f"INI: ", f0) - print(f"Sum(f0) = ", np.sum(f0)) - - # After time + dt: - dt = 1.0e7 - ft = func_solver_eigenval(natom, te0, ne0, time + dt, f0, table) - - print(f"time_end = ", time + dt) - print(f"NEI:", ft) - print(f"Sum(ft) = ", np.sum(ft)) - print(f"EI :", table.equilibrium_state(T_e=te0)) - print("End Test.\n") - - assert np.isclose(np.sum(ft), 1), "np.sum(ft) is not approximately 1" - assert np.isclose(np.sum(f0), 1), "np.sum(f0) is not approximately 1" - assert np.allclose(ft, table.equilibrium_state(T_e=te0)) - - -@pytest.mark.xfail -def test_reachequlibrium_state_multisteps(natom=8): - """ - Starting the random initial distribution, the charge states will reach - to equilibrium cases after a long time (multiple steps). - In this test, we set the ionization and recombination rates at - Te0=2.0e6 K and plasma density ne0=1.0e+7. A random charge states - distribution will be finally closed to equilibrium distribution at - 2.0e6K. - """ - # - # Initial conditions, set plasma temperature, density and dt - # - element_symbol = atomic.atomic_symbol(int(natom)) - te0 = 1.0e6 # unit: K - ne0 = 1.0e8 # unit: cm^-3 - - # Start from any ionizaiont states, e.g., Te = 4.0d4 K, - time = 0 - table = EigenData2(element=natom) - f0 = table.equilibrium_state(T_e=4.0e4) - - print("START test_reachequlibrium_state_multisteps:") - print(f"time_sta = ", time) - print(f"INI: ", f0) - print(f"Sum(f0) = ", np.sum(f0)) - - # After time + dt: - dt = 100000.0 # unit: second - - # Enter the time loop: - for it in range(100): - ft = func_solver_eigenval(natom, te0, ne0, time + dt, f0, table) - f0 = np.copy(ft) - time = time + dt - - print(f"time_end = ", time + dt) - print(f"NEI:", ft) - print(f"Sum(ft) = ", np.sum(ft)) - - assert np.isclose(np.sum(ft), 1) - - print(f"EI :", table.equilibrium_state(T_e=te0)) - print("End Test.\n") - - -# Temporarily test only lighter elements due to Travis CI delays - -# @pytest.mark.parametrize('atomic_numb', np.arange(1, 27)) -@pytest.mark.xfail -@pytest.mark.parametrize("atomic_numb", np.arange(1, 10)) -def test_element_range(atomic_numb): - """ - Function test_element_range: - This function is used to test element including Hydrogen to Iron. - """ - try: - element_symbol = atomic.atomic_symbol(int(atomic_numb)) - eigen = EigenData2(element=element_symbol) - except Exception as exc: - raise Exception(f"Problem with atomic number={atomic_numb}.") from exc - - eigen = 0 # attempt to clear up memory From 672ec0015806a369d9f8124d43b49d39c9ef07b5 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Thu, 30 Apr 2020 20:37:51 -0400 Subject: [PATCH 09/49] Copy NEI class and tests from NEI-modeling/NEI Co-authored-by: Chengcai Shen Co-authored-by: Marcus Dupont --- plasmapy_nei/nei/nei.py | 1565 +++++++++++++++++++++++++++- plasmapy_nei/nei/tests/test_nei.py | 297 ++++++ 2 files changed, 1861 insertions(+), 1 deletion(-) diff --git a/plasmapy_nei/nei/nei.py b/plasmapy_nei/nei/nei.py index 0111d02..a40705b 100644 --- a/plasmapy_nei/nei/nei.py +++ b/plasmapy_nei/nei/nei.py @@ -3,5 +3,1568 @@ __all__ = ["NEI"] -class NEI: +import numpy as np +import matplotlib.pyplot as plt +from typing import Union, Optional, List, Dict, Callable +import astropy.units as u +import plasmapy as pl +from scipy import interpolate, optimize +from .eigenvaluetable import EigenData2 +from .ionization_states import IonizationStates +import warnings +from nei.physics import * + +# TODO: Allow this to keep track of velocity and position too, and +# eventually to have density and temperature be able to be functions of +# position. (and more complicated expressions for density and +# temperature too) + +# TODO: Expand Simulation docstring + + +class NEIError(Exception): pass + + +class Simulation: + """ + Results from a non-equilibrium ionization simulation. + + Parameters + ---------- + initial: ~IonizationStates + The `IonizationStates` instance representing the ionization + states of different elements and plasma properties as the + initial conditions. + + n_init: ~astropy.units.quantity + The initial number density scaling factor. + + T_e_init: ~astropy.units.quantity + The initial electron temperature. + + max_steps: int + The maximum number of time steps that the simulation can take + before stopping. + + time_start: ~astropy.units.Quantity + The time at the start of the simulation. + + """ + def __init__( + self, + initial: IonizationStates, + n_init: u.Quantity, + T_e_init: u.Quantity, + max_steps: int, + time_start: u.Quantity, + ): + + self._elements = initial.elements + self._abundances = initial.abundances + self._max_steps = max_steps + + self._nstates = {elem: pl.atomic.atomic_number(elem) + 1 + for elem in self.elements} + + self._ionic_fractions = { + elem: np.full((max_steps + 1, self.nstates[elem]), np.nan, + dtype=np.float64) + for elem in self.elements + } + + self._number_densities = { + elem: np.full((max_steps + 1, self.nstates[elem]), np.nan, + dtype=np.float64) * u.cm ** -3 + for elem in self.elements + } + + self._n_elem = { + elem: np.full(max_steps + 1, np.nan) * u.cm ** -3 + for elem in self.elements + } + + self._n_e = np.full(max_steps + 1, np.nan) * u.cm ** -3 + self._T_e = np.full(max_steps + 1, np.nan) * u.K + self._time = np.full(max_steps + 1, np.nan) * u.s + + self._index = 0 + + self._assign( + new_time=time_start, + new_ionfracs=initial.ionic_fractions, + new_n = n_init, + new_T_e = T_e_init, + ) + + def _assign( + self, + new_time: u.Quantity, + new_ionfracs: Dict[str, np.ndarray], + new_n: u.Quantity, + new_T_e: u.Quantity, + ): + """ + Store results from a time step of a non-equilibrium ionization + time advance in the `~nei.classes.NEI` class. + + Parameters + ---------- + new_time: ~astropy.units.Quantity + The time associated with this time step. + + new_ionfracs: dict + The new ionization fractions for this time step. The keys + of this `dict` are the atomic symbols of the elements being + tracked, and with the corresponding value being an + `~numpy.ndarray` representing the ionic fractions. Each + element's array must have a length of the atomic number plus + one, and be normalized to one with all values between zero + and one. + + new_n: ~astropy.units.Quantity + The new number density scaling factor for this time step. + The number densities of each ionic species will be the + product of this scaling factor, the element's abundance, and + the ionic fraction given in `new_ionfracs`. + + new_T_e: ~astropy.units.Quantity + The new electron temperature. + + """ + + try: + index = self._index + elements = self.elements + self._time[index] = new_time + self._T_e[index] = new_T_e + + for elem in elements: + self._ionic_fractions[elem][index, :] = new_ionfracs[elem][:] + + # Calculate elemental and ionic number densities + n_elem = {elem: new_n * self.abundances[elem] for elem in elements} + number_densities = { + elem: n_elem[elem] * new_ionfracs[elem] + for elem in elements + } + + # Calculate the electron number density + n_e = 0.0 * u.cm ** -3 + for elem in elements: + integer_charges = np.linspace( + 0, self.nstates[elem]-1, self.nstates[elem]) + n_e += np.sum(number_densities[elem] * integer_charges) + + # Assign densities + self._n_e[index] = n_e + for elem in elements: + self._n_elem[elem][index] = n_elem[elem] + self._number_densities[elem][index, :] = number_densities[elem] + + except Exception as exc: + raise NEIError( + f"Unable to assign parameters to Simulation instance " + f"for index {index} at time = {new_time}. The " + f"parameters are new_n = {new_n}, new_T_e = {new_T_e}, " + f"and new_ionic_fractions = {new_ionfracs}." + ) from exc + finally: + self._index += 1 + + def _cleanup(self): + """ + Clean up this class after the simulation is complete. + + This method removes the excess elements from each array that + did not end up getting used for a time step in the simulation + and sets the `last_step` attribute. + + """ + nsteps = self._index + + self._n_e = self._n_e[0:nsteps] + self._T_e = self._T_e[0:nsteps] + self._time = self._time[0:nsteps] + + for element in self.elements: + self._ionic_fractions[element] = \ + self._ionic_fractions[element][0:nsteps, :] + self._number_densities[element] = \ + self._number_densities[element][0:nsteps, :] + + self._last_step = nsteps - 1 + + self._index = None + + @property + def max_steps(self) -> int: + """ + The maximum number of time steps allowed for this simulation. + """ + return self._max_steps + + @property + def last_step(self) -> int: + """ + The time index of the last step. + """ + return self._last_step + + @property + def nstates(self) -> Dict[str, int]: + """ + Return the dictionary containing atomic symbols as keys and the + number of ionic species for the corresponding element as the + value. + """ + return self._nstates + + @property + def elements(self) -> List[str]: + """The elements modeled by this simulation.""" + return self._elements + + @property + def abundances(self) -> Dict[str, float]: + """ + The relative elemental abundances of the elements modeled in + this simulation. + + The keys are the atomic symbols and the values are a `float` + representing that element's elemental abundance. + + """ + return self._abundances + + @property + def ionic_fractions(self) -> Dict[str, np.ndarray]: + """ + Return the ionic fractions over the course of the simulation. + + The keys of this dictionary are atomic symbols. The values are + 2D arrays where the first index refers to the time step and the + second index refers to the integer charge. + + """ + return self._ionic_fractions + + @property + def number_densities(self) -> Dict[str, u.Quantity]: + """ + Return the number densities over the course of the simulation. + + The keys of `number_densities` are atomic symbols. The values + are 2D arrays with units of number density where the first index + refers to the time step and the second index is the integer + charge. + + """ + return self._number_densities + + @property + def n_elem(self) -> Dict[str, u.Quantity]: + """ + The number densities of each element over the course of the + simulation. + + The keys of `n_elem` are atomic symbols. The values are 1D + arrays with units of number density where the index refers to + the time step. + + """ + return self._n_elem + + @property + def n_e(self) -> u.Quantity: + """ + The electron number density over the course of the simulation in + units of number density. + + The index of this array corresponds to the time step. + """ + return self._n_e + + @property + def T_e(self) -> u.Quantity: + """ + The electron temperature over the course of the simulation in + kelvin. + + The index of this array corresponds to the time step. + """ + return self._T_e + + @property + def time(self) -> u.Quantity: + """ + The time for each time step over the course of the simulation + in units of seconds. + """ + return self._time + + +class NEI: + r""" + Perform and analyze a non-equilibrium ionization simulation. + + Parameters + ---------- + inputs + + T_e: `~astropy.units.Quantity` or `callable` + The electron temperature, which may be a constant, an array of + temperatures corresponding to the times in `time_input`, or a + function that yields the temperature as a function of time. + + n: `~astropy.units.Quantity` or `callable` + The number density multiplicative factor. The number density of + each element will be `n` times the abundance given in + `abundances`. For example, if `abundance['H'] = 1`, then this + will correspond to the number density of hydrogen (including + neutral hydrogen and protons). This factor may be a constant, + an array of number densities over time, or a function that + yields a number density as a function of time. + + time_input: ~astropy.units.Quantity, optional + An array containing the times associated with `n` and `T_e` in + units of time. + + time_start: ~astropy.units.Quantity, optional + The start time for the simulation. If density and/or + temperature are given by arrays, then this argument must be + greater than `time_input[0]`. If this argument is not supplied, + then `time_start` defaults to `time_input[0]` (if given) and + zero seconds otherwise. + + time_max: ~astropy.units.Quantity + The maximum time for the simulation. If density and/or + temperature are given by arrays, then this argument must be less + than `time_input[-1]`. + + max_steps: `int` + The maximum number of time steps to be taken during a + simulation. + + dt: `~astropy.units.Quantity` + The time step. If `adapt_dt` is `False`, then `dt` is the time + step for the whole simulation. + + dt_max: `~astropy.units.Quantity` + The maximum time step to be used with an adaptive time step. + + dt_min: `~astropy.units.Quantity` + The minimum time step to be used with an adaptive time step. + + adapt_dt: `bool` + If `True`, change the time step based on the characteristic + ionization and recombination time scales and change in + temperature. Not yet implemented. + + safety_factor: `float` or `int` + A multiplicative factor to multiply by the time step when + `adapt_dt` is `True`. Lower values improve accuracy, whereas + higher values reduce computational time. Not yet implemented. + + tol: float + The absolute tolerance to be used in comparing ionic fractions. + + verbose: bool, optional + A flag stating whether or not to print out information for every + time step. Setting `verbose` to `True` is useful for testing. + Defaults to `False`. + + abundances: dict + + Examples + -------- + + >>> import numpy as np + >>> import astropy.units as u + + >>> inputs = {'H': [0.9, 0.1], 'He': [0.9, 0.099, 0.001]} + >>> abund = {'H': 1, 'He': 0.085} + >>> n = np.array([1e9, 1e8]) * u.cm ** -3 + >>> T_e = np.array([10000, 40000]) * u.K + >>> time = np.array([0, 300]) * u.s + >>> dt = 0.25 * u.s + + The initial conditions can be accessed using the initial attribute. + + >>> sim = NEI(inputs=inputs, abundances=abund, n=n, T_e=T_e, time_input=time, adapt_dt=False, dt=dt) + + After having inputted all of the necessary information, we can run + the simulation. + + >>> results = sim.simulate() + + The initial results are stored in the `initial` attribute. + + >>> sim.initial.ionic_fractions['H'] + array([0.9, 0.1]) + + The final results can be access with the `final` attribute. + + >>> sim.final.ionic_fractions['H'] + array([0.16665179, 0.83334821]) + >>> sim.final.ionic_fractions['He'] + array([0.88685261, 0.11218358, 0.00096381]) + >>> sim.final.T_e + + + Both `initial` and `final` are instances of the `IonizationStates` + class. + + Notes + ----- + The ionization and recombination rates are from Chianti version + 8.x. These rates include radiative and dielectronic recombination. + Photoionization is not included. + + """ + + def __init__( + self, + inputs, + abundances: Union[Dict, str] = None, + T_e: Union[Callable, u.Quantity] = None, + n: Union[Callable, u.Quantity] = None, + time_input: u.Quantity = None, + time_start: u.Quantity = None, + time_max: u.Quantity = None, + max_steps: Union[int, np.integer] = 10000, + tol: Union[int, float] = 1e-15, + dt: u.Quantity = None, + dt_max: u.Quantity = np.inf * u.s, + dt_min: u.Quantity = 0 * u.s, + adapt_dt: bool = None, + safety_factor: Union[int, float] = 1, + verbose: bool = False, + ): + + try: + + self.time_input = time_input + self.time_start = time_start + self.time_max = time_max + self.T_e_input = T_e + self.n_input = n + self.max_steps = max_steps + self.dt_input = dt + + if self.dt_input is None: + self._dt = self.time_max / max_steps + else: + self._dt = self.dt_input + + self.dt_min = dt_min + self.dt_max = dt_max + self.adapt_dt = adapt_dt + self.safety_factor = safety_factor + self.verbose = verbose + + T_e_init = self.electron_temperature(self.time_start) + n_init = self.hydrogen_number_density(self.time_start) + + self.initial = IonizationStates( + inputs=inputs, + abundances=abundances, + T_e=T_e_init, + n_H=n_init, # TODO: Update n_H in IonizationState(s) + tol = tol, + ) + + self.tol = tol + self.elements = self.initial.elements + + if 'H' not in self.elements: + raise NEIError("Must have H in elements") + + self.abundances = self.initial.abundances + + self._EigenDataDict = { + element: EigenData2(element) for element in self.elements + } + + if self.T_e_input is not None and not isinstance(inputs, dict): + for element in self.initial.elements: + self.initial.ionic_fractions[element] = \ + self.EigenDataDict[element].equilibrium_state(T_e_init.value) + + self._temperature_grid = \ + self._EigenDataDict[self.elements[0]].temperature_grid + + self._get_temperature_index = \ + self._EigenDataDict[self.elements[0]]._get_temperature_index + + except Exception: + raise NEIError( + f"Unable to create NEI instance for:\n" + f" inputs = {inputs}\n" + f" abundances = {abundances}\n" + f" T_e = {T_e}\n" + f" n = {n}\n" + f" time_input = {time_input}\n" + f" time_start = {time_start}\n" + f" time_max = {time_max}\n" + f" max_steps = {max_steps}\n" + ) + + def equil_ionic_fractions( + self, + T_e: u.Quantity = None, + time: u.Quantity = None, + ) -> Dict[str, np.ndarray]: + """ + Return the equilibrium ionic fractions for a temperature or at + a given time. + + Parameters + ---------- + T_e: ~astropy.units.Quantity, optional + The electron temperature in units that can be converted to + kelvin. + + time: ~astropy.units.Quantity, optional + The time in units that can be converted to seconds. + + Returns + ------- + equil_ionfracs: dict + The equilibrium ionic fractions for the elements contained + within this class + + Notes + ----- + Only one of `T_e` and `time` may be included as an argument. If + neither `T_e` or `time` is provided and the temperature for the + simulation is given by a constant, the this method will assume + that `T_e` is the temperature of the simulation. + + """ + + if T_e is not None and time is not None: + raise NEIError("Only one of T_e and time may be an argument.") + + if T_e is None and time is None: + if self.T_e_input.isscalar: + T_e = self.T_e_input + else: + raise NEIError + + try: + T_e = T_e.to(u.K) if T_e is not None else None + time = time.to(u.s) if time is not None else None + except Exception as exc: + raise NEIError("Invalid input to equilibrium_ionic_fractions.") + + if time is not None: + T_e = self.electron_temperature(time) + + if not T_e.isscalar: + raise NEIError("Need scalar input for equil_ionic_fractions.") + + equil_ionfracs = {} + for element in self.elements: + equil_ionfracs[element] = \ + self.EigenDataDict[element].equilibrium_state(T_e.value) + + return equil_ionfracs + + @property + def elements(self) -> List[str]: + return self._elements + + @elements.setter + def elements(self, elements): + # TODO: Update this + self._elements = elements + + @property + def abundances(self) -> Dict[str, Union[float, int]]: + """Return the abundances.""" + return self._abundances + + @abundances.setter + def abundances(self, abund: Dict[Union[str, int], Union[float, int]]): + + # TODO: Update initial, etc. when abundances is updated. The + # checks within IonizationStates will also be checks for + + # TODO: Update initial and other attributes when abundances is + # updated. + + self._abundances = abund + + @property + def tol(self) -> float: + """ + The tolerance for comparisons between different ionization + states. + """ + return self._tol + + @tol.setter + def tol(self, value: Union[float, int]): + try: + value = float(value) + except Exception as exc: + raise TypeError("Invalid tolerance.") + if not 0 <= value < 1: + raise ValueError("Need 0 <= tol < 1.") + self._tol = value + + @property + def time_input(self) -> Optional[u.Quantity]: + return self._time_input + + @time_input.setter + def time_input(self, times: Optional[u.Quantity]): + if times is None: + self._time_input = None + elif isinstance(times, u.Quantity): + if times.isscalar: + raise ValueError("time_input must be an array.") + try: + times = times.to(u.s) + except u.UnitConversionError: + raise u.UnitsError( + "time_input must have units of seconds.") from None + if not np.all(times[1:] > times[:-1]): + raise ValueError("time_input must monotonically increase.") + self._time_input = times + else: + raise TypeError("Invalid time_input.") + + @property + def time_start(self) -> u.Quantity: + """The start time of the simulation.""" + return self._time_start + + @time_start.setter + def time_start(self, time: u.Quantity): + if time is None: + self._time_start = 0.0 * u.s + elif isinstance(time, u.Quantity): + if not time.isscalar: + raise ValueError("time_start must be a scalar") + try: + time = time.to(u.s) + except u.UnitConversionError: + raise u.UnitsError( + "time_start must have units of seconds") from None + if hasattr(self, '_time_max') \ + and self._time_max is not None \ + and self._time_max<=time: + raise ValueError("Need time_start < time_max.") + if self.time_input is not None and \ + self.time_input.min() > time: + raise ValueError( + "time_start must be less than min(time_input)") + self._time_start = time + else: + raise TypeError("Invalid time_start.") from None + + @property + def time_max(self) -> u.Quantity: + """The maximum time allowed for the simulation.""" + return self._time_max + + @time_max.setter + def time_max(self, time: u.Quantity): + if time is None: + self._time_max = self.time_input[-1] \ + if self.time_input is not None else np.inf * u.s + elif isinstance(time, u.Quantity): + if not time.isscalar: + raise ValueError("time_max must be a scalar") + try: + time = time.to(u.s) + except u.UnitConversionError: + raise u.UnitsError( + "time_max must have units of seconds") from None + if hasattr(self, '_time_start') and self._time_start is not None \ + and self._time_start >= time: + raise ValueError("time_max must be greater than time_start") + self._time_max = time + else: + raise TypeError("Invalid time_max.") from None + + @property + def adapt_dt(self) -> Optional[bool]: + """ + Return `True` if the time step is set to be adaptive, `False` + if the time step is set to not be adapted, and `None` if this + attribute was not set. + """ + return self._adapt_dt + + @adapt_dt.setter + def adapt_dt(self, choice: Optional[bool]): + if choice is None: + self._adapt_dt = True if self.dt_input is None else False + elif choice is True or choice is False: + self._adapt_dt = choice + else: + raise TypeError("Invalid value for adapt_dt") + + @property + def dt_input(self) -> Optional[u.Quantity]: + """Return the inputted time step.""" + return self._dt_input + + @dt_input.setter + def dt_input(self, dt: Optional[u.Quantity]): + if dt is None: + self._dt_input = None + elif isinstance(dt, u.Quantity): + try: + dt = dt.to(u.s) + if dt > 0 * u.s: + self._dt_input = dt + except (AttributeError, u.UnitConversionError): + raise NEIError("Invalid dt.") + + @property + def dt_min(self) -> u.Quantity: + return self._dt_min + + @dt_min.setter + def dt_min(self, value: u.Quantity): + if not isinstance(value, u.Quantity): + raise TypeError("dt_min must be a Quantity.") + try: + value = value.to(u.s) + except u.UnitConversionError as exc: + raise u.UnitConversionError("Invalid units for dt_min.") + if hasattr(self, '_dt_input') and self.dt_input is not None and self.dt_input < value: + raise ValueError( + "dt_min cannot exceed the inputted time step.") + if hasattr(self, '_dt_max') and self.dt_max < value: + raise ValueError( + "dt_min cannot exceed dt_max.") + self._dt_min = value + + @property + def dt_max(self) -> u.Quantity: + return self._dt_max + + @dt_max.setter + def dt_max(self, value: u.Quantity): + if not isinstance(value, u.Quantity): + raise TypeError("dt_max must be a Quantity.") + try: + value = value.to(u.s) + except u.UnitConversionError as exc: + raise u.UnitConversionError("Invalid units for dt_max.") + if hasattr(self, '_dt_input') and self.dt_input is not None and self.dt_input > value: + raise ValueError( + "dt_max cannot be less the inputted time step.") + if hasattr(self, '_dt_min') and self.dt_min > value: + raise ValueError( + "dt_min cannot exceed dt_max.") + self._dt_max = value + + @property + def safety_factor(self): + """ + The multiplicative factor that the time step is to be multiplied + by when using an adaptive time step. + """ + return self._safety_factor + + @safety_factor.setter + def safety_factor(self, value): + if not isinstance(value, (float, np.float64, np.integer, int)): + raise TypeError + if 1e-3 <= value <= 1e3: + self._safety_factor = value + else: + raise NEIError("Invalid safety factor.") + + @property + def verbose(self) -> bool: + """ + Return `True` if verbose output during a simulation is + requested, and `False` otherwise. + """ + return self._verbose + + @verbose.setter + def verbose(self, choice: bool): + if choice is True or choice is False: + self._verbose = choice + else: + raise TypeError("Invalid choice for verbose.") + + @u.quantity_input + def in_time_interval(self, time: u.s, buffer: u.s = 1e-9 * u.s): + """ + Return `True` if the `time` is between `time_start - buffer` and + `time_max + buffer` , and `False` otherwise. + + Raises + ------ + TypeError + If `time` or `buffer` is not a `~astropy.units.Quantity` + + `~astropy.units.UnitsError` + If `time` or `buffer` is not in units of time. + + """ + return self.time_start - buffer <= time <= self.time_max + buffer + + @property + def max_steps(self) -> int: + """ + The maximum number of steps that a simulation will be allowed + to take. + """ + return self._max_steps + + @max_steps.setter + def max_steps(self, n: int): + if isinstance(n, (int, np.integer)) and 0 < n <= 1000000: + self._max_steps = n + else: + raise TypeError( + "max_steps must be an integer with 0 < max_steps <= " + "1000000") + + @property + def T_e_input(self) -> Union[u.Quantity, Callable]: + """ + The temperature input. + """ + return self._T_e_input + + @T_e_input.setter + def T_e_input(self, T_e: Optional[Union[Callable, u.Quantity]]): + """Set the input electron temperature.""" + if isinstance(T_e, u.Quantity): + try: + T_e = T_e.to(u.K, equivalencies=u.temperature_energy()) + except u.UnitConversionError: + raise u.UnitsError("Invalid electron temperature.") from None + if T_e.isscalar: + self._T_e_input = T_e + self._electron_temperature = lambda time: T_e + else: + if self._time_input is None: + raise TypeError( + "Must define time_input prior to T_e for an array.") + time_input = self.time_input + if len(time_input) != len(T_e): + raise ValueError("len(T_e) not equal to len(time_input).") + f = interpolate.interp1d( + time_input.value, + T_e.value, + bounds_error=False, + fill_value="extrapolate", + ) + self._electron_temperature = lambda time: f(time.value) * u.K + self._T_e_input = T_e + elif callable(T_e): + if self.time_start is not None: + try: + T_e(self.time_start).to(u.K) + T_e(self.time_max).to(u.K) + except Exception: + raise ValueError("Invalid electron temperature function.") + self._T_e_input = T_e + self._electron_temperature = T_e + elif T_e is None: + self._electron_temperature = lambda: None + else: + raise TypeError("Invalid T_e") + + def electron_temperature(self, time: u.Quantity) -> u.Quantity: + try: + if not self.in_time_interval(time): + warnings.warn( + f"{time} is not in the simulation time interval:" + f"[{self.time_start}, {self.time_max}]. " + f"May be extrapolating temperature.") + T_e = self._electron_temperature(time.to(u.s)) + if np.isnan(T_e) or np.isinf(T_e) or T_e < 0 * u.K: + raise NEIError(f"T_e = {T_e} at time = {time}.") + return T_e + except Exception as exc: + raise NEIError( + f"Unable to calculate a valid electron temperature " + f"for time {time}") from exc + + @property + def n_input(self) -> u.Quantity: + """The number density factor input.""" + if 'H' in self.elements: + return self._n_input + else: + raise ValueError + + @n_input.setter + def n_input(self, n: u.Quantity): + if isinstance(n, u.Quantity): + try: + n = n.to(u.cm ** -3) + except u.UnitConversionError: + raise u.UnitsError("Invalid hydrogen density.") + if n.isscalar: + self._n_input = n + self.hydrogen_number_density = lambda time: n + else: + if self._time_input is None: + raise TypeError( + "Must define time_input prior to n for an array.") + time_input = self.time_input + if len(time_input) != len(n): + raise ValueError("len(n) is not equal to len(time_input).") + f = interpolate.interp1d( + time_input.value, + n.value, + bounds_error=False, + fill_value="extrapolate", + ) + self._hydrogen_number_density = \ + lambda time: f(time.value) * u.cm ** -3 + self._n_input = n + elif callable(n): + if self.time_start is not None: + try: + n(self.time_start).to(u.cm ** -3) + n(self.time_max).to(u.cm ** -3) + except Exception: + raise ValueError("Invalid number density function.") + self._n_input = n + self._hydrogen_number_density = n + elif n is None: + self._hydrogen_number_density = lambda: None + else: + raise TypeError("Invalid n.") + + def hydrogen_number_density(self, time: u.Quantity) -> u.Quantity: + try: + time = time.to(u.s) + except (AttributeError, u.UnitsError): + raise NEIError("Invalid time in hydrogen_density") + return self._hydrogen_number_density(time) + + @property + def EigenDataDict(self) -> Dict[str, EigenData2]: + """ + Return a `dict` containing `~nei.class + """ + return self._EigenDataDict + + @property + def initial(self) -> IonizationStates: + """ + Return the ionization states of the plasma at the beginning of + the simulation. + """ + return self._initial + + @initial.setter + def initial(self, initial_states: Optional[IonizationStates]): + if isinstance(initial_states, IonizationStates): + self._initial = initial_states + self._elements = initial_states.elements + elif initial_states is None: + self._ionstates = None + else: + raise TypeError("Expecting an IonizationStates instance.") + + @property + def results(self) -> Simulation: + """ + Return the `~nei.classes.Simulation` class instance that + corresponds to the simulation results. + + """ + try: + return self._results + except Exception: + raise AttributeError("The simulation has not yet been performed.") + + @property + def final(self) -> IonizationStates: + """ + Return the ionization states of the plasma at the end of the + simulation. + """ + try: + return self._final + except AttributeError: + raise NEIError("The simulation has not yet been performed.") from None + + def _initialize_simulation(self): + + self._results = Simulation( + initial=self.initial, + n_init=self.hydrogen_number_density(self.time_start), + T_e_init=self.electron_temperature(self.time_start), + max_steps=self.max_steps, + time_start=self.time_start, + ) + self._old_time = self.time_start.to(u.s) + self._new_time = self.time_start.to(u.s) + + def simulate(self) -> Simulation: + """ + Perform a non-equilibrium ionization simulation. + + Returns + ------- + results: ~nei.classes.Simulation + The results from the simulation (which are also stored in + the `results` attribute of the `~nei.classes.NEI` instance + this method was called from. + + """ + + self._initialize_simulation() + + for step in range(self.max_steps): + try: + self.set_timestep() + self.time_advance() + except StopIteration: + break + except Exception as exc: + raise NEIError(f"Unable to complete simulation.") from exc + + self._finalize_simulation() + + # Is there a way to use the inspect package or something similar + # to only return self.results if it is in an expression where + + return self.results + + def _finalize_simulation(self): + self._results._cleanup() + + final_ionfracs = { + element: self.results.ionic_fractions[element][-1, :] + for element in self.elements + } + + self._final = IonizationStates( + inputs=final_ionfracs, + abundances=self.abundances, + n_H=np.sum(self.results.number_densities['H'][-1, :]), # modify this later?, + T_e=self.results.T_e[-1], + tol=1e-6, + ) + + if not np.isclose(self.time_max/u.s, self.results.time[-1]/u.s): + warnings.warn( + f"The simulation ended at {self.results.time[-1]}, " + f"which is prior to time_max = {self.time_max}." + ) + + def _set_adaptive_timestep(self): + """Adapt the time step.""" + + t = self._new_time if hasattr(self, '_new_time') else self.t_start + + # We need to guess the timestep in order to narrow down what the + # timestep should be. If we are in the middle of a simulation, + # we can use the old timestep as a reasonable guess. If we are + # simulation, then we can either use the inputted timestep or + # estimate it from other inputs. + + dt_guess = self._dt if self._dt \ + else self._dt_input if self._dt_input \ + else self.time_max / self.max_steps + + # Make sure that dt_guess does not lead to a time that is out + # of the domain. + + dt_guess = dt_guess if t + dt_guess <= self.time_max - t else self.time_max - t + + # The temperature may start out exactly at the boundary of a + # bin, so we check what bin it is in just slightly after to + # figure out which temperature bin the plasma is entering. + + T = self.electron_temperature(t + 1e-9 * dt_guess) + + # Find the boundaries to the temperature bin. + + index = self._get_temperature_index(T.to(u.K).value) + T_nearby = np.array(self._temperature_grid[index-1:index+2]) * u.K + T_boundary = (T_nearby[0:-1] + T_nearby[1:]) / 2 + + # In order to use Brent's method, we must bound the root's + # location. Functions may change sharply or slowly, so we test + # different times that are logarithmically spaced to find the + # first one that is outside of the boundary. + + dt_spread = np.geomspace(1e-9 * dt_guess.value, (self.time_max - t).value, num=100) * u.s + time_spread = t + dt_spread + T_spread = [self.electron_temperature(time) for time in time_spread] + in_range = [T_boundary[0] <= temp <= T_boundary[1] for temp in T_spread] + + # If all of the remaining temperatures are in the same bin, then + # the temperature will be roughly constant for the rest of the + # simulation. Take one final long time step, unless it exceeds + # dt_max. + + if all(in_range): + new_dt = self.time_max - t + self._dt = new_dt if new_dt <= self.dt_max else self.dt_max + return + + # Otherwise, we need to find the first index in the spread that + # corresponds to a temperature outside of the temperature bin + # for this time step. + + first_false_index = in_range.index(False) + + # We need to figure out if the temperature is dropping so that + # it crosses the low temperature boundary of the bin, or if it + # is rising so that it crosses the high temperature of the bin. + + T_first_outside = self.electron_temperature(time_spread[first_false_index]) + + if T_first_outside >= T_boundary[1]: + boundary_index = 1 + elif T_first_outside <= T_boundary[0]: + boundary_index = 0 + + # Select the values for the time step in the spread just before + # and after the temperature leaves the temperature bin as bounds + # for the root finding method. + + dt_bounds = (dt_spread[first_false_index-1:first_false_index+1]).value + + # Define a function for the difference between the temperature + # and the temperature boundary as a function of the value of the + # time step. + + T_val = lambda dtval: \ + (self.electron_temperature(t + dtval*u.s) - T_boundary[boundary_index]).value + + # Next we find the root. This method should succeed as long as + # the root is bracketed by dt_bounds. Because astropy.units is + # not fully compatible with SciPy, we temporarily drop units and + # then reattach them. + + try: + new_dt = optimize.brentq( + T_val, + *dt_bounds, + xtol=1e-14, + maxiter=1000, + disp=True, + ) * u.s + except Exception as exc: + raise NEIError(f"Unable to find new dt at t = {t}") from exc + else: + if np.isnan(new_dt.value): + raise NEIError(f"new_dt = {new_dt}") + + # Enforce that the time step is in the interval [dt_min, dt_max]. + + if new_dt < self.dt_min: + new_dt = self.dt_min + elif new_dt > self.dt_max: + new_dt = self.dt_max + + # Store the time step as a private attribute so that it can be + # used in the time advance. + + self._dt = new_dt.to(u.s) + + def set_timestep(self, dt: u.Quantity = None): + """ + Set the time step for the next non-equilibrium ionization time + advance. + + Parameters + ---------- + dt: ~astropy.units.Quantity, optional + The time step to be used for the next time advance. + + Notes + ----- + If `dt` is not `None`, then the time step will be set to `dt`. + + If `dt` is not set and the `adapt_dt` attribute of an + `~nei.classes.NEI` instance is `True`, then this method will + calculate the time step corresponding to how long it will be + until the temperature rises or drops into the next temperature + bin. If this time step is between `dtmin` and `dtmax`, then + + If `dt` is not set and the `adapt_dt` attribute is `False`, then + this method will set the time step as what was inputted to the + `~nei.classes.NEI` class upon instantiation in the `dt` argument + or through the `~nei.classes.NEI` class's `dt_input` attribute. + + Raises + ------ + ~nei.classes.NEIError + If the time step cannot be set, for example if the `dt` + argument is invalid or the time step cannot be adapted. + + """ + + if dt is not None: + # Allow the time step to set as an argument to this method. + try: + dt = dt.to(u.s) + except Exception as exc: + raise NEIError(f"{dt} is not a valid time step.") from exc + finally: + self._dt = dt + elif self.adapt_dt: + try: + self._set_adaptive_timestep() + except Exception as exc: + raise NEIError("Unable to adapt the time step.") from exc + elif self.dt_input is not None: + self._dt = self.dt_input + else: + raise NEIError("Unable to set the time step.") + + self._old_time = self._new_time + self._new_time = self._old_time + self._dt + + if self._new_time > self.time_max: + self._new_time = self.time_max + self._dt = self._new_time - self._old_time + + def time_advance(self): + """ + Advance the simulation by one time step. + """ + # TODO: Expand docstring and include equations! + + # TODO: Fully implement units into this. + + step = self.results._index + T_e = self.results.T_e[step - 1].value + n_e = self.results.n_e[step - 1].value # set average + dt = self._dt.value + + if self.verbose: + print( + f"step={step} T_e={T_e} n_e={n_e} dt={dt}" + ) + + new_ionic_fractions = {} + + try: + for elem in self.elements: + nstates = self.results.nstates[elem] + f0 = self.results._ionic_fractions[elem][self.results._index - 1, :] + + evals = self.EigenDataDict[elem].eigenvalues(T_e=T_e) + evect = self.EigenDataDict[elem].eigenvectors(T_e=T_e) + evect_inverse = self.EigenDataDict[elem].eigenvector_inverses(T_e=T_e) + + diagonal_evals = np.zeros((nstates, nstates), dtype=np.float64) + for ii in range(0, nstates): + diagonal_evals[ii, ii] = np.exp(evals[ii] * dt * n_e) + + matrix_1 = np.dot(diagonal_evals, evect) + matrix_2 = np.dot(evect_inverse, matrix_1) + + ft = np.dot(f0, matrix_2) + + # Due to truncation errors in the solutions in the + # eigenvalues and eigenvectors, there is a chance that + # very slightly negative ionic fractions will arise. + # These are not natural and will make the code grumpy. + # For these reasons, the ionic fractions will be very + # slightly unnormalized. We set negative ionic + # fractions to zero and renormalize. + + ft[np.where(ft < 0.0)] = 0.0 + new_ionic_fractions[elem] = ft / np.sum(ft) + + except Exception as exc: + raise NEIError(f"Unable to do time advance for {elem}") from exc + else: + + new_time = self.results.time[self.results._index-1] + self._dt + self.results._assign( + new_time=new_time, + new_ionfracs=new_ionic_fractions, + new_T_e=self.electron_temperature(new_time), + new_n=self.hydrogen_number_density(new_time), + ) + + if new_time >= self.time_max or np.isclose(new_time.value, self.time_max.value): + raise StopIteration + + def save(self, filename: str = "nei.h5"): + """ + Save the `~nei.classes.NEI` instance to an HDF5 file. Not + implemented. + """ + raise NotImplementedError + + def visual(self, element): + """ + Returns an atomic object used for plotting protocols + + Parameter + ------ + element: str, + The elemental symbol of the particle in question (i.e. 'H') + + Returns + ------ + Class object + """ + + plot_obj = Visualize(element, self.results) + + return plot_obj + + def index_to_time(self, index): + """ + Returns the time value or array given the index/indices + + Parameters + ------ + index: array-like + A value or array of values representing the index of + the time array created by the simulation + + Returns + ------ + get_time: astropy.units.Quantity + The time value associated with index input(s) + """ + + return self.results.time[index] + + def time_to_index(self, time): + """ + Returns the closest index value or array for the given time(s) + + Parameters + ------ + time: array-like, + A value or array of values representing the values of + the time array created by the simulation + + Returns + ------ + index: int or array-like, + The index value associated with the time input(s) + """ + index = (np.abs(self.results.time.value - time)).argmin() + + return index + +class Visualize(NEI): + """ + Store plotting results from the simulation + """ + def __init__(self, element, results): + self.element = element + self._results = results + + def index_to_time(self, index): + """ + Inherits the index_to_time method of the NEI class + """ + return super(Visualize, self).index_to_time(index) + + + def ionicfrac_evol_plot(self,x_axis, ion='all'): + """ + Creates a plot of the ionic fraction time evolution of element inputs + + Paramaters + ------ + element: string, + The elemental symbal of the atom (i.e. 'H'), + ion: array-like, dtype=int + The repective integer charge of the atomic particle (i.e. 0 or [0,1]) + + x_axis: ~astropy.units.Quantity: array-like , + The xaxis to plot the ionic fraction evolution over. + Can only be distance or time. + + """ + #Check if element input is a string + if not isinstance(self.element, str): + raise TypeError('The element input must be a string') + + + #Check to see if x_axis units are of time or length + if isinstance(x_axis, u.Quantity): + try: + x = x_axis.to(u.s) + xlabel = 'Time' + except Exception: + x = x_axis.to(u.R_sun) + xlabel = 'Distance' + else: + raise TypeError('Invalid x-axis units. Must be units of length or time.') + + + if ion == 'all': + charge_states = pl.atomic.atomic_number(self.element) + 1 + else: + #Ensure ion input is array of integers + charge_states = np.array(ion, dtype=np.int16) + + linsetyles = ['--','-.',':','-'] + + if ion =='all': + for nstate in range(charge_states): + ionic_frac = self.results.ionic_fractions[self.element][:,nstate] + plt.plot(x.value,ionic_frac, linestyle= linsetyles[nstate % len(linsetyles)], label=f'{self.element}+{nstate}') + plt.xlabel(f'{xlabel} ({x.unit})') + plt.ylabel('Ionic Fraction') + plt.title('Ionic Fraction Evolution of {}'.format(self.element)) + plt.legend(loc='center left') + else: + if charge_states.size > 1: + for nstate in charge_states: + ionic_frac = self.results.ionic_fractions[self.element][:,nstate] + plt.plot(x.value,ionic_frac, linestyle= linsetyles[nstate % len(linsetyles)], label=f'{self.element}+{nstate}') + plt.xlabel(f'{xlabel} ({x.unit})') + plt.ylabel('Ionic Fraction') + plt.title('Ionic Fraction Evolution of {}'.format(self.element)) + plt.legend() + #plt.show() + else: + ionic_frac = self.results.ionic_fractions[self.element][:,charge_states] + plt.plot(x.value,ionic_frac) + plt.xlabel(f'{xlabel} ({x.unit})') + plt.ylabel('Ionic Fraction') + plt.title('Ionic Fraction Evolution of %s$^{%i+}$'%(self.element,ion)) + #plt.show() + + def ionicfrac_bar_plot(self, time_index): + """ + Creates a bar plot of the ion fraction change at a particular time index + + Parameters + ------ + time_index: int, + The particular time index at which to collect the various ion fractiom + change + """ + + if not isinstance(self.element, str): + raise TypeError('The element input must be a string') + + ion = pl.atomic.atomic_number(self.element) + + charge_states = np.linspace(0, ion, ion+1, dtype=np.int16) + + width=1.0 + + fig, ax = plt.subplots() + if isinstance(time_index, (list, np.ndarray)): + + alpha = 1.0 + colors = ['blue', 'red'] + + #Color index counter + color_idx = 1 + + for idx in time_index: + + #Toggle between zero and one for colors array + color_idx ^= 1 + + ax.bar(charge_states, self.results.ionic_fractions[self.element][idx,:], alpha=alpha, \ + width=width, color=colors[color_idx], + label='Time:{time:.{number}f}'.format(time=self.index_to_time(idx), number=1)) + alpha -= 0.2 + ax.set_xticks(charge_states-width/2.0) + ax.set_xticklabels(charge_states) + ax.set_title(f'{self.element}') + + ax.set_xlabel('Charge State') + ax.set_ylabel('Ionic Fraction') + + ax.legend(loc='best') + #plt.show() + + else: + ax.bar(x, self.results.ionic_fractions[self.element][time_index,:], alpha=1.0, width=width) + ax.set_xticks(charge_states-width/2.0) + ax.set_xticklabels(charge_states) + ax.set_title(f'{self.element}') + ax.set_xlabel('Charge State') + ax.set_ylabel('Ionic Fraction') + #plt.show() + + def rh_density_plot(self, gamma, mach, ion='None'): + """ + Creates a plot of the Rankine-Huguniot jump relation for the + density of our element + + Parameters + ------ + gamma: float, + The specific heats ratio of the system + mach: float, + The mach number of the scenario + ion: int, + The ionic integer charge of the element in question + """ + + #Instantiate the MHD class + mhd = shocks.MHD() + + nstates = pl.atomic.atomic_number(self.element) + 1 + + if ion == 'None': + + for charge in range(nstates): + + post_rho = mhd.rh_density(self.results.number_densities[self.element].value[0, charge], gamma, mach) + + plt.semilogy(post_rho, label=f'{self.element}{charge}+') + + plt.legend() + #plt.show() + + else: + + if not isinstance(ion, int): + raise TypeError('Please make sure that your charge value is an integer') + elif ion > nstates: + raise ValueError('The ionic charge input is greater than allowed for this element') + + + post_rho = mhd.rh_density(init_dens = self.results.number_densities[self.element].value[0, ion], gamma=gamma, mach=mach) + + plt.semilogy(post_rho) + plt.title('Density Shock Transition') + plt.xlabel('Mach Number') + plt.ylabel('Number Density') + #plt.show() + + def rh_temp_plot(self, gamma, mach): + """ + Creates a plot of the Rankine-Huguniot jump relation for the + density of our element + + Parameters + ------ + gamma: float, + The specific heats ratio of the system + mach: float, + The mach number of the scenario + """ + + #Instantiate the MHD class + mhd = shocks.MHD() + + post_temp = mhd.rh_temp(self.results.T_e.value[0], gamma, mach) + + plt.semilogy(mach, post_temp) + plt.title('Temperature Shock Transition') + plt.xlabel('Mach Number') + plt.ylabel('Log Temperature (K)') + #plt.show() diff --git a/plasmapy_nei/nei/tests/test_nei.py b/plasmapy_nei/nei/tests/test_nei.py index 1ae83bc..d5851e0 100644 --- a/plasmapy_nei/nei/tests/test_nei.py +++ b/plasmapy_nei/nei/tests/test_nei.py @@ -4,3 +4,300 @@ def test_import(): """Minimal test that importing this works.""" from plasmapy_nei.nei import NEI + + +import astropy.units as u +from ..ionization_states import IonizationStates, particle_symbol +from ..nei import NEI +from ..eigenvaluetable import EigenData2 +import numpy as np +import pytest + +inputs_dict = {'H': [0.9, 0.1], 'He': [0.5, 0.3, 0.2]} +abundances = {'H': 1, 'He': 0.1} + +time_array = np.array([0, 800]) * u.s +T_e_array = np.array([4e4, 6e4]) * u.K +n_array = np.array([1e9, 5e8]) * u.cm ** -3 + + +tests = { + + 'basic': { + 'inputs': inputs_dict, + 'abundances': abundances, + 'T_e': T_e_array, + 'n': n_array, + 'time_input': time_array, + 'time_start': 0 * u.s, + 'time_max': 800 * u.s, + 'max_steps': 1, + 'dt': 800 * u.s, + 'adapt_dt': False, + 'verbose': True, + }, + + 'T_e constant': { + 'inputs': inputs_dict, + 'abundances': abundances, + 'T_e': 1 * u.MK, + 'n': n_array, + 'time_input': time_array, + 'time_start': 0 * u.s, + 'time_max': 800 * u.s, + 'dt': 100 * u.s, + 'max_steps': 2, + 'adapt_dt': False, + 'verbose': True, + }, + + 'n_e constant': { + 'inputs': inputs_dict, + 'abundances': abundances, + 'T_e': T_e_array, + 'n': 1e9 * u.cm ** -3, + 'time_input': time_array, + 'time_start': 0 * u.s, + 'time_max': 800 * u.s, + 'max_steps': 2, + 'adapt_dt': False, + 'dt': 100 * u.s, + 'verbose': True, + }, + + 'T_e function': { + 'inputs': inputs_dict, + 'abundances': abundances, + 'T_e': lambda time: 1e4 * (1 + time/u.s) * u.K, + 'n': 1e15 * u.cm **-3, + 'time_max': 800 * u.s, + 'max_steps': 2, + 'dt': 100 * u.s, + 'adapt_dt': False, + 'verbose': True, + }, + + 'n function': { + 'inputs': inputs_dict, + 'abundances': abundances, + 'T_e': 6e4 * u.K, + 'n': lambda time: 1e9 * (1 + time/u.s) * u.cm ** -3, + 'time_start': 0 * u.s, + 'time_max': 800 * u.s, + 'adapt_dt': False, + 'dt': 200 * u.s, + 'verbose': True, + 'max_steps': 4 + }, + + 'equil test cool': { + 'inputs': ['H', 'He', 'N'], + 'abundances': {'H': 1, 'He': 0.1, 'C': 1e-4, 'N': 1e-4, 'O': 1e-4, 'Fe': 1e-4}, + 'T_e': 10001.0 * u.K, + 'n': 1e13 * u.cm ** -3, + 'time_max': 2e6 * u.s, + 'tol': 1e-9, + 'adapt_dt': False, + 'dt': 5e5 * u.s, + 'max_steps': 4, + 'verbose': True, + }, + + 'equil test hot': { + 'inputs': ['H', 'He', 'C'], + 'abundances': {'H': 1, 'He': 0.1, 'C': 1e-4, 'N': 1e-4, 'O': 1e-4, 'Fe': 1e-4, 'S': 2e-6}, + 'T_e': 7e6 * u.K, + 'n': 1e9 * u.cm ** -3, + 'time_max': 1e8 * u.s, + 'dt': 5e7 * u.s, + 'max_steps': 3, + 'adapt_dt': False, + 'verbose': True, + }, + + 'equil test start far out of equil': { + 'inputs': { + 'H': [0.99, 0.01], + 'He': [0.5, 0.0, 0.5], + 'O': [0.2, 0, 0.2, 0, 0.2, 0, 0.2, 0, 0.2], + }, + 'abundances': {'H': 1, 'He': 0.1, 'O': 1e-4}, + 'T_e': 3e6 * u.K, + 'n': 1e9 * u.cm ** -3, + 'dt': 1e6 * u.s, + 'time_start': 0 * u.s, + 'time_max': 1e6 * u.s, + 'adapt_dt': False, + 'verbose': True, + 'max_steps': 2 + }, + + 'adapt dt': { + 'inputs': ['H', 'He'], + 'abundances': {'H': 1, 'He': 0.1}, + 'T_e': lambda t: u.K * (1e6 + 1.3e4*np.sin(t.value)), + 'n': 1e10 * u.cm ** -3, + 'max_steps': 300, + 'time_start': 0 * u.s, + 'time_max': 2*np.pi * u.s, + 'adapt_dt': True, + }, +} + +test_names = list(tests.keys()) + + +class TestNEI: + + @classmethod + def setup_class(cls): + cls.instances = {} + + @pytest.mark.parametrize('test_name', test_names) + def test_instantiate(self, test_name): + try: + instance = NEI(**tests[test_name]) + self.instances[test_name] = instance + except Exception as exc: + raise Exception(f"Problem with test {test_name}") from exc + + @pytest.mark.parametrize('test_name', test_names) + def test_time_start(self, test_name): + instance = self.instances[test_name] + if 'time_start' in tests[test_name].keys(): + assert tests[test_name]['time_start'] == instance.time_start + elif 'time_input' in tests[test_name].keys(): + assert tests[test_name]['time_input'].min() == instance.time_start + else: + assert instance.time_start == 0 * u.s + + @pytest.mark.parametrize('test_name', test_names) + def test_time_max(self, test_name): + instance = self.instances[test_name] + if 'time_max' in tests[test_name].keys(): + assert tests[test_name]['time_max'] == instance.time_max + + @pytest.mark.parametrize('test_name', test_names) + def test_initial_type(self, test_name): + instance = self.instances[test_name] + assert isinstance(instance.initial, IonizationStates) + + @pytest.mark.parametrize('test_name', test_names) + def test_n_input(self, test_name): + actual = self.instances[test_name].n_input + expected = tests[test_name]['n'] + if isinstance(expected, u.Quantity) and not expected.isscalar: + assert all(expected == actual) + else: + assert expected == actual + + @pytest.mark.parametrize('test_name', test_names) + def test_T_e_input(self, test_name): + actual = self.instances[test_name].T_e_input + expected = tests[test_name]['T_e'] + if isinstance(expected, u.Quantity) and not expected.isscalar: + assert all(expected == actual) + else: + assert expected == actual + + @pytest.mark.parametrize('test_name', test_names) + def test_electron_temperature(self, test_name): + instance = self.instances[test_name] + T_e_input = instance.T_e_input + if isinstance(T_e_input, u.Quantity): + if T_e_input.isscalar: + assert instance.electron_temperature(instance.time_start) == T_e_input + assert instance.electron_temperature(instance.time_max) == T_e_input + else: + for time, T_e in zip(instance.time_input, T_e_input): + try: + T_e_func = instance.electron_temperature(time) + except Exception: + raise ValueError("Unable to find T_e from electron_temperature") + + assert np.isclose(T_e.value, T_e_func.value) + if callable(T_e_input): + assert instance.T_e_input(instance.time_start) == \ + instance.electron_temperature(instance.time_start) + + @pytest.mark.parametrize( + 'test_name', + [test_name for test_name in test_names if isinstance(tests[test_name]['inputs'], dict)], + ) + def test_initial_ionfracs(self, test_name): + original_inputs = tests[test_name]['inputs'] + original_elements = original_inputs.keys() + + for element in original_elements: + assert np.allclose( + original_inputs[element], + self.instances[test_name].initial.ionic_fractions[particle_symbol(element)] + ) + + @pytest.mark.parametrize('test_name', test_names) + def test_simulate(self, test_name): + try: + self.instances[test_name].simulate() + except Exception as exc: + raise ValueError(f"Unable to simulate for test: {test_name}") from exc + + @pytest.mark.parametrize('test_name', test_names) + def test_simulation_end(self, test_name): + instance = self.instances[test_name] + time = instance.results.time + end_time = time[-1] + max_steps = instance.max_steps + + if np.isnan(end_time.value): + raise Exception('End time is NaN.') + + got_to_max_steps = len(time) == instance.max_steps + 1 + got_to_time_max = np.isclose(time[-1].value, instance.time_max.value) + + if time.isscalar or len(time) == 1: + raise Exception(f"The only element in results.time is {time}") + + if not got_to_max_steps and not got_to_time_max: + print(f"time = {time}") + print(f"max_steps = {max_steps}") + print(f'time_max = {instance.time_max}') + raise Exception("Problem with end time.") + + @pytest.mark.parametrize( + 'test_name', + [test_name for test_name in test_names if 'equil' in test_name] + ) + def test_equilibration(self, test_name): + """ + Test that equilibration works. + """ + instance = self.instances[test_name] + T_e = instance.T_e_input + assert isinstance(T_e, u.Quantity) and T_e.isscalar, \ + "This test can only be used for cases where T_e is constant." + equil_dict = instance.equil_ionic_fractions(T_e) + for element in instance.elements: + assert np.allclose( + equil_dict[element], instance.results.ionic_fractions[element][-1, :] + ) + + @pytest.mark.parametrize('test_name', test_names) + def test_initial_results(self, test_name): + initial = self.instances[test_name].initial + results = self.instances[test_name].results + assert initial.elements == results.elements + assert initial.abundances == results.abundances + for elem in initial.elements: + assert np.allclose(results.ionic_fractions[elem][0, :], initial.ionic_fractions[elem]) + + @pytest.mark.parametrize('test_name', test_names) + def test_final_results(self, test_name): + #initial = self.instances[test_name].initial + final = self.instances[test_name].final + results = self.instances[test_name].results + assert final.elements == results.elements + assert final.abundances == results.abundances + for elem in final.elements: + assert np.allclose( + results.ionic_fractions[elem][-1, :], final.ionic_fractions[elem] + ) From 2d246f3b7f7b48cd9ae3853decb3d36422dbc4b2 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Thu, 30 Apr 2020 21:17:17 -0400 Subject: [PATCH 10/49] Remove nearly empty test file in eigen --- plasmapy_nei/eigen/tests/test_eigen.py | 6 ------ 1 file changed, 6 deletions(-) delete mode 100644 plasmapy_nei/eigen/tests/test_eigen.py diff --git a/plasmapy_nei/eigen/tests/test_eigen.py b/plasmapy_nei/eigen/tests/test_eigen.py deleted file mode 100644 index 1d2d124..0000000 --- a/plasmapy_nei/eigen/tests/test_eigen.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Tests for eigentables""" - - -def test_import(): - """Test that the subpackage can be imported.""" - import plasmapy_nei.eigen From a889b883a7300a1f882f5ac5182b20efee21a129 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Thu, 30 Apr 2020 21:17:40 -0400 Subject: [PATCH 11/49] Remove example subpackage left over from templating --- plasmapy_nei/example_subpkg/__init__.py | 5 ----- plasmapy_nei/example_subpkg/data/.gitignore | 0 plasmapy_nei/example_subpkg/tests/__init__.py | 0 3 files changed, 5 deletions(-) delete mode 100644 plasmapy_nei/example_subpkg/__init__.py delete mode 100644 plasmapy_nei/example_subpkg/data/.gitignore delete mode 100644 plasmapy_nei/example_subpkg/tests/__init__.py diff --git a/plasmapy_nei/example_subpkg/__init__.py b/plasmapy_nei/example_subpkg/__init__.py deleted file mode 100644 index 621b0a7..0000000 --- a/plasmapy_nei/example_subpkg/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -""" -This is the docstring for the examplesubpkg package. Normally you would -have whatever.py files in this directory implementing some modules, but this -is just an example sub-package, so it doesn't actually do anything. -""" diff --git a/plasmapy_nei/example_subpkg/data/.gitignore b/plasmapy_nei/example_subpkg/data/.gitignore deleted file mode 100644 index e69de29..0000000 diff --git a/plasmapy_nei/example_subpkg/tests/__init__.py b/plasmapy_nei/example_subpkg/tests/__init__.py deleted file mode 100644 index e69de29..0000000 From 7df4d97624a8be3cbab119f2b31f66c72e14f0e7 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Thu, 30 Apr 2020 21:20:26 -0400 Subject: [PATCH 12/49] Adapt NEI class & tests to this repository --- plasmapy_nei/__init__.py | 6 +- plasmapy_nei/nei/__init__.py | 2 +- plasmapy_nei/nei/nei.py | 521 +++++++++-------------------- plasmapy_nei/nei/tests/test_nei.py | 321 +++++++++--------- 4 files changed, 324 insertions(+), 526 deletions(-) diff --git a/plasmapy_nei/__init__.py b/plasmapy_nei/__init__.py index 6d5f8ec..4c37b70 100644 --- a/plasmapy_nei/__init__.py +++ b/plasmapy_nei/__init__.py @@ -4,15 +4,13 @@ try: from .version import __version__ + del version except Exception as exc: warnings.warn("Unable to import __version__") -else: - del version finally: del warnings from . import eigen from . import nei -# Then you can be explicit to control what ends up in the namespace, -__all__ = ["eigen"] +__all__ = ["eigen", "nei"] diff --git a/plasmapy_nei/nei/__init__.py b/plasmapy_nei/nei/__init__.py index da06f7b..239364c 100644 --- a/plasmapy_nei/nei/__init__.py +++ b/plasmapy_nei/nei/__init__.py @@ -1,3 +1,3 @@ """Classes that perform non-equilibrium ionization modeling""" -from plasmapy_nei.nei.nei import NEI +from plasmapy_nei.nei.nei import NEI, NEIError diff --git a/plasmapy_nei/nei/nei.py b/plasmapy_nei/nei/nei.py index a40705b..43bb4ea 100644 --- a/plasmapy_nei/nei/nei.py +++ b/plasmapy_nei/nei/nei.py @@ -1,18 +1,18 @@ """Contains classes to represent non-equilibrium ionization simulations.""" -__all__ = ["NEI"] +__all__ = ["NEI", "NEIError"] import numpy as np -import matplotlib.pyplot as plt from typing import Union, Optional, List, Dict, Callable import astropy.units as u import plasmapy as pl from scipy import interpolate, optimize -from .eigenvaluetable import EigenData2 -from .ionization_states import IonizationStates +from plasmapy_nei.eigen import EigenData +from plasmapy.atomic import IonizationStates import warnings -from nei.physics import * + + # TODO: Allow this to keep track of velocity and position too, and # eventually to have density and temperature be able to be functions of @@ -22,6 +22,18 @@ # TODO: Expand Simulation docstring +# TODO: Include the methods in the original Visualize class which is a +# subclass of NEI in the NEI-modeling/NEI repo. These were deleted +# temporarily to make it possible to get the NEI class itself +# adapted into this package. + + +# TODO: In this file and test_nei.py, there are a few places with +# initial.ionic_fractions.keys(), where initial is an instance +# of IonizationStates. This workaround exists because I forgot +# to put in an `elements` attribute in IonizationStates, and +# should be corrected. + class NEIError(Exception): pass @@ -32,7 +44,7 @@ class Simulation: Parameters ---------- - initial: ~IonizationStates + initial: plasmapy.atomic.IonizationStates The `IonizationStates` instance representing the ionization states of different elements and plasma properties as the initial conditions. @@ -51,37 +63,37 @@ class Simulation: The time at the start of the simulation. """ + def __init__( - self, - initial: IonizationStates, - n_init: u.Quantity, - T_e_init: u.Quantity, - max_steps: int, - time_start: u.Quantity, + self, + initial: IonizationStates, + n_init: u.Quantity, + T_e_init: u.Quantity, + max_steps: int, + time_start: u.Quantity, ): - self._elements = initial.elements + self._elements = list(initial.ionic_fractions.keys()) self._abundances = initial.abundances self._max_steps = max_steps - self._nstates = {elem: pl.atomic.atomic_number(elem) + 1 - for elem in self.elements} + self._nstates = { + elem: pl.atomic.atomic_number(elem) + 1 for elem in self.elements + } self._ionic_fractions = { - elem: np.full((max_steps + 1, self.nstates[elem]), np.nan, - dtype=np.float64) + elem: np.full((max_steps + 1, self.nstates[elem]), np.nan, dtype=np.float64) for elem in self.elements } self._number_densities = { - elem: np.full((max_steps + 1, self.nstates[elem]), np.nan, - dtype=np.float64) * u.cm ** -3 + elem: np.full((max_steps + 1, self.nstates[elem]), np.nan, dtype=np.float64) + * u.cm ** -3 for elem in self.elements } self._n_elem = { - elem: np.full(max_steps + 1, np.nan) * u.cm ** -3 - for elem in self.elements + elem: np.full(max_steps + 1, np.nan) * u.cm ** -3 for elem in self.elements } self._n_e = np.full(max_steps + 1, np.nan) * u.cm ** -3 @@ -93,20 +105,20 @@ def __init__( self._assign( new_time=time_start, new_ionfracs=initial.ionic_fractions, - new_n = n_init, - new_T_e = T_e_init, + new_n=n_init, + new_T_e=T_e_init, ) def _assign( - self, - new_time: u.Quantity, - new_ionfracs: Dict[str, np.ndarray], - new_n: u.Quantity, - new_T_e: u.Quantity, + self, + new_time: u.Quantity, + new_ionfracs: Dict[str, np.ndarray], + new_n: u.Quantity, + new_T_e: u.Quantity, ): """ Store results from a time step of a non-equilibrium ionization - time advance in the `~nei.classes.NEI` class. + time advance in the `~plasmapy_nei.classes.NEI` class. Parameters ---------- @@ -145,15 +157,15 @@ def _assign( # Calculate elemental and ionic number densities n_elem = {elem: new_n * self.abundances[elem] for elem in elements} number_densities = { - elem: n_elem[elem] * new_ionfracs[elem] - for elem in elements + elem: n_elem[elem] * new_ionfracs[elem] for elem in elements } # Calculate the electron number density n_e = 0.0 * u.cm ** -3 for elem in elements: integer_charges = np.linspace( - 0, self.nstates[elem]-1, self.nstates[elem]) + 0, self.nstates[elem] - 1, self.nstates[elem] + ) n_e += np.sum(number_densities[elem] * integer_charges) # Assign densities @@ -188,10 +200,10 @@ def _cleanup(self): self._time = self._time[0:nsteps] for element in self.elements: - self._ionic_fractions[element] = \ - self._ionic_fractions[element][0:nsteps, :] - self._number_densities[element] = \ - self._number_densities[element][0:nsteps, :] + self._ionic_fractions[element] = self._ionic_fractions[element][0:nsteps, :] + self._number_densities[element] = self._number_densities[element][ + 0:nsteps, : + ] self._last_step = nsteps - 1 @@ -424,22 +436,22 @@ class NEI: """ def __init__( - self, - inputs, - abundances: Union[Dict, str] = None, - T_e: Union[Callable, u.Quantity] = None, - n: Union[Callable, u.Quantity] = None, - time_input: u.Quantity = None, - time_start: u.Quantity = None, - time_max: u.Quantity = None, - max_steps: Union[int, np.integer] = 10000, - tol: Union[int, float] = 1e-15, - dt: u.Quantity = None, - dt_max: u.Quantity = np.inf * u.s, - dt_min: u.Quantity = 0 * u.s, - adapt_dt: bool = None, - safety_factor: Union[int, float] = 1, - verbose: bool = False, + self, + inputs, + abundances: Union[Dict, str] = None, + T_e: Union[Callable, u.Quantity] = None, + n: Union[Callable, u.Quantity] = None, + time_input: u.Quantity = None, + time_start: u.Quantity = None, + time_max: u.Quantity = None, + max_steps: Union[int, np.integer] = 10000, + tol: Union[int, float] = 1e-15, + dt: u.Quantity = None, + dt_max: u.Quantity = np.inf * u.s, + dt_min: u.Quantity = 0 * u.s, + adapt_dt: bool = None, + safety_factor: Union[int, float] = 1, + verbose: bool = False, ): try: @@ -470,32 +482,38 @@ def __init__( inputs=inputs, abundances=abundances, T_e=T_e_init, - n_H=n_init, # TODO: Update n_H in IonizationState(s) - tol = tol, + n=n_init, + tol=tol, ) self.tol = tol - self.elements = self.initial.elements - if 'H' not in self.elements: + # TODO: Update IonizationStates in PlasmaPy to have elements attribute + + self.elements = list(self.initial.ionic_fractions.keys()) + + if "H" not in self.elements: raise NEIError("Must have H in elements") self.abundances = self.initial.abundances self._EigenDataDict = { - element: EigenData2(element) for element in self.elements + element: EigenData(element) for element in self.elements } if self.T_e_input is not None and not isinstance(inputs, dict): - for element in self.initial.elements: - self.initial.ionic_fractions[element] = \ - self.EigenDataDict[element].equilibrium_state(T_e_init.value) + for element in self.initial.ionic_fractions.keys(): + self.initial.ionic_fractions[element] = self.EigenDataDict[ + element + ].equilibrium_state(T_e_init.value) - self._temperature_grid = \ - self._EigenDataDict[self.elements[0]].temperature_grid + self._temperature_grid = self._EigenDataDict[ + self.elements[0] + ].temperature_grid - self._get_temperature_index = \ - self._EigenDataDict[self.elements[0]]._get_temperature_index + self._get_temperature_index = self._EigenDataDict[ + self.elements[0] + ]._get_temperature_index except Exception: raise NEIError( @@ -511,9 +529,7 @@ def __init__( ) def equil_ionic_fractions( - self, - T_e: u.Quantity = None, - time: u.Quantity = None, + self, T_e: u.Quantity = None, time: u.Quantity = None, ) -> Dict[str, np.ndarray]: """ Return the equilibrium ionic fractions for a temperature or at @@ -556,7 +572,7 @@ def equil_ionic_fractions( T_e = T_e.to(u.K) if T_e is not None else None time = time.to(u.s) if time is not None else None except Exception as exc: - raise NEIError("Invalid input to equilibrium_ionic_fractions.") + raise NEIError("Invalid input to equilibrium_ionic_fractions.") from exc if time is not None: T_e = self.electron_temperature(time) @@ -566,8 +582,9 @@ def equil_ionic_fractions( equil_ionfracs = {} for element in self.elements: - equil_ionfracs[element] = \ - self.EigenDataDict[element].equilibrium_state(T_e.value) + equil_ionfracs[element] = self.EigenDataDict[element].equilibrium_state( + T_e.value + ) return equil_ionfracs @@ -628,8 +645,7 @@ def time_input(self, times: Optional[u.Quantity]): try: times = times.to(u.s) except u.UnitConversionError: - raise u.UnitsError( - "time_input must have units of seconds.") from None + raise u.UnitsError("time_input must have units of seconds.") from None if not np.all(times[1:] > times[:-1]): raise ValueError("time_input must monotonically increase.") self._time_input = times @@ -651,16 +667,15 @@ def time_start(self, time: u.Quantity): try: time = time.to(u.s) except u.UnitConversionError: - raise u.UnitsError( - "time_start must have units of seconds") from None - if hasattr(self, '_time_max') \ - and self._time_max is not None \ - and self._time_max<=time: + raise u.UnitsError("time_start must have units of seconds") from None + if ( + hasattr(self, "_time_max") + and self._time_max is not None + and self._time_max <= time + ): raise ValueError("Need time_start < time_max.") - if self.time_input is not None and \ - self.time_input.min() > time: - raise ValueError( - "time_start must be less than min(time_input)") + if self.time_input is not None and self.time_input.min() > time: + raise ValueError("time_start must be less than min(time_input)") self._time_start = time else: raise TypeError("Invalid time_start.") from None @@ -673,18 +688,21 @@ def time_max(self) -> u.Quantity: @time_max.setter def time_max(self, time: u.Quantity): if time is None: - self._time_max = self.time_input[-1] \ - if self.time_input is not None else np.inf * u.s + self._time_max = ( + self.time_input[-1] if self.time_input is not None else np.inf * u.s + ) elif isinstance(time, u.Quantity): if not time.isscalar: raise ValueError("time_max must be a scalar") try: time = time.to(u.s) except u.UnitConversionError: - raise u.UnitsError( - "time_max must have units of seconds") from None - if hasattr(self, '_time_start') and self._time_start is not None \ - and self._time_start >= time: + raise u.UnitsError("time_max must have units of seconds") from None + if ( + hasattr(self, "_time_start") + and self._time_start is not None + and self._time_start >= time + ): raise ValueError("time_max must be greater than time_start") self._time_max = time else: @@ -737,12 +755,14 @@ def dt_min(self, value: u.Quantity): value = value.to(u.s) except u.UnitConversionError as exc: raise u.UnitConversionError("Invalid units for dt_min.") - if hasattr(self, '_dt_input') and self.dt_input is not None and self.dt_input < value: - raise ValueError( - "dt_min cannot exceed the inputted time step.") - if hasattr(self, '_dt_max') and self.dt_max < value: - raise ValueError( - "dt_min cannot exceed dt_max.") + if ( + hasattr(self, "_dt_input") + and self.dt_input is not None + and self.dt_input < value + ): + raise ValueError("dt_min cannot exceed the inputted time step.") + if hasattr(self, "_dt_max") and self.dt_max < value: + raise ValueError("dt_min cannot exceed dt_max.") self._dt_min = value @property @@ -757,12 +777,14 @@ def dt_max(self, value: u.Quantity): value = value.to(u.s) except u.UnitConversionError as exc: raise u.UnitConversionError("Invalid units for dt_max.") - if hasattr(self, '_dt_input') and self.dt_input is not None and self.dt_input > value: - raise ValueError( - "dt_max cannot be less the inputted time step.") - if hasattr(self, '_dt_min') and self.dt_min > value: - raise ValueError( - "dt_min cannot exceed dt_max.") + if ( + hasattr(self, "_dt_input") + and self.dt_input is not None + and self.dt_input > value + ): + raise ValueError("dt_max cannot be less the inputted time step.") + if hasattr(self, "_dt_min") and self.dt_min > value: + raise ValueError("dt_min cannot exceed dt_max.") self._dt_max = value @property @@ -828,8 +850,8 @@ def max_steps(self, n: int): self._max_steps = n else: raise TypeError( - "max_steps must be an integer with 0 < max_steps <= " - "1000000") + "max_steps must be an integer with 0 < max_steps <= " "1000000" + ) @property def T_e_input(self) -> Union[u.Quantity, Callable]: @@ -851,8 +873,7 @@ def T_e_input(self, T_e: Optional[Union[Callable, u.Quantity]]): self._electron_temperature = lambda time: T_e else: if self._time_input is None: - raise TypeError( - "Must define time_input prior to T_e for an array.") + raise TypeError("Must define time_input prior to T_e for an array.") time_input = self.time_input if len(time_input) != len(T_e): raise ValueError("len(T_e) not equal to len(time_input).") @@ -884,20 +905,21 @@ def electron_temperature(self, time: u.Quantity) -> u.Quantity: warnings.warn( f"{time} is not in the simulation time interval:" f"[{self.time_start}, {self.time_max}]. " - f"May be extrapolating temperature.") + f"May be extrapolating temperature." + ) T_e = self._electron_temperature(time.to(u.s)) if np.isnan(T_e) or np.isinf(T_e) or T_e < 0 * u.K: raise NEIError(f"T_e = {T_e} at time = {time}.") return T_e except Exception as exc: raise NEIError( - f"Unable to calculate a valid electron temperature " - f"for time {time}") from exc + f"Unable to calculate a valid electron temperature " f"for time {time}" + ) from exc @property def n_input(self) -> u.Quantity: """The number density factor input.""" - if 'H' in self.elements: + if "H" in self.elements: return self._n_input else: raise ValueError @@ -914,8 +936,7 @@ def n_input(self, n: u.Quantity): self.hydrogen_number_density = lambda time: n else: if self._time_input is None: - raise TypeError( - "Must define time_input prior to n for an array.") + raise TypeError("Must define time_input prior to n for an array.") time_input = self.time_input if len(time_input) != len(n): raise ValueError("len(n) is not equal to len(time_input).") @@ -925,8 +946,7 @@ def n_input(self, n: u.Quantity): bounds_error=False, fill_value="extrapolate", ) - self._hydrogen_number_density = \ - lambda time: f(time.value) * u.cm ** -3 + self._hydrogen_number_density = lambda time: f(time.value) * u.cm ** -3 self._n_input = n elif callable(n): if self.time_start is not None: @@ -950,7 +970,7 @@ def hydrogen_number_density(self, time: u.Quantity) -> u.Quantity: return self._hydrogen_number_density(time) @property - def EigenDataDict(self) -> Dict[str, EigenData2]: + def EigenDataDict(self) -> Dict[str, EigenData]: """ Return a `dict` containing `~nei.class """ @@ -968,7 +988,7 @@ def initial(self) -> IonizationStates: def initial(self, initial_states: Optional[IonizationStates]): if isinstance(initial_states, IonizationStates): self._initial = initial_states - self._elements = initial_states.elements + self._elements = initial_states.ionic_fractions.keys() # TODO IonizationStates elif initial_states is None: self._ionstates = None else: @@ -1051,12 +1071,14 @@ def _finalize_simulation(self): self._final = IonizationStates( inputs=final_ionfracs, abundances=self.abundances, - n_H=np.sum(self.results.number_densities['H'][-1, :]), # modify this later?, + n=np.sum( + self.results.number_densities["H"][-1, :] + ), # modify this later?, T_e=self.results.T_e[-1], tol=1e-6, ) - if not np.isclose(self.time_max/u.s, self.results.time[-1]/u.s): + if not np.isclose(self.time_max / u.s, self.results.time[-1] / u.s): warnings.warn( f"The simulation ended at {self.results.time[-1]}, " f"which is prior to time_max = {self.time_max}." @@ -1065,7 +1087,7 @@ def _finalize_simulation(self): def _set_adaptive_timestep(self): """Adapt the time step.""" - t = self._new_time if hasattr(self, '_new_time') else self.t_start + t = self._new_time if hasattr(self, "_new_time") else self.t_start # We need to guess the timestep in order to narrow down what the # timestep should be. If we are in the middle of a simulation, @@ -1073,9 +1095,13 @@ def _set_adaptive_timestep(self): # simulation, then we can either use the inputted timestep or # estimate it from other inputs. - dt_guess = self._dt if self._dt \ - else self._dt_input if self._dt_input \ + dt_guess = ( + self._dt + if self._dt + else self._dt_input + if self._dt_input else self.time_max / self.max_steps + ) # Make sure that dt_guess does not lead to a time that is out # of the domain. @@ -1091,7 +1117,7 @@ def _set_adaptive_timestep(self): # Find the boundaries to the temperature bin. index = self._get_temperature_index(T.to(u.K).value) - T_nearby = np.array(self._temperature_grid[index-1:index+2]) * u.K + T_nearby = np.array(self._temperature_grid[index - 1 : index + 2]) * u.K T_boundary = (T_nearby[0:-1] + T_nearby[1:]) / 2 # In order to use Brent's method, we must bound the root's @@ -1099,7 +1125,10 @@ def _set_adaptive_timestep(self): # different times that are logarithmically spaced to find the # first one that is outside of the boundary. - dt_spread = np.geomspace(1e-9 * dt_guess.value, (self.time_max - t).value, num=100) * u.s + dt_spread = ( + np.geomspace(1e-9 * dt_guess.value, (self.time_max - t).value, num=100) + * u.s + ) time_spread = t + dt_spread T_spread = [self.electron_temperature(time) for time in time_spread] in_range = [T_boundary[0] <= temp <= T_boundary[1] for temp in T_spread] @@ -1135,14 +1164,15 @@ def _set_adaptive_timestep(self): # and after the temperature leaves the temperature bin as bounds # for the root finding method. - dt_bounds = (dt_spread[first_false_index-1:first_false_index+1]).value + dt_bounds = (dt_spread[first_false_index - 1 : first_false_index + 1]).value # Define a function for the difference between the temperature # and the temperature boundary as a function of the value of the # time step. - T_val = lambda dtval: \ - (self.electron_temperature(t + dtval*u.s) - T_boundary[boundary_index]).value + T_val = lambda dtval: ( + self.electron_temperature(t + dtval * u.s) - T_boundary[boundary_index] + ).value # Next we find the root. This method should succeed as long as # the root is bracketed by dt_bounds. Because astropy.units is @@ -1150,13 +1180,10 @@ def _set_adaptive_timestep(self): # then reattach them. try: - new_dt = optimize.brentq( - T_val, - *dt_bounds, - xtol=1e-14, - maxiter=1000, - disp=True, - ) * u.s + new_dt = ( + optimize.brentq(T_val, *dt_bounds, xtol=1e-14, maxiter=1000, disp=True,) + * u.s + ) except Exception as exc: raise NEIError(f"Unable to find new dt at t = {t}") from exc else: @@ -1247,9 +1274,7 @@ def time_advance(self): dt = self._dt.value if self.verbose: - print( - f"step={step} T_e={T_e} n_e={n_e} dt={dt}" - ) + print(f"step={step} T_e={T_e} n_e={n_e} dt={dt}") new_ionic_fractions = {} @@ -1286,7 +1311,7 @@ def time_advance(self): raise NEIError(f"Unable to do time advance for {elem}") from exc else: - new_time = self.results.time[self.results._index-1] + self._dt + new_time = self.results.time[self.results._index - 1] + self._dt self.results._assign( new_time=new_time, new_ionfracs=new_ionic_fractions, @@ -1304,24 +1329,6 @@ def save(self, filename: str = "nei.h5"): """ raise NotImplementedError - def visual(self, element): - """ - Returns an atomic object used for plotting protocols - - Parameter - ------ - element: str, - The elemental symbol of the particle in question (i.e. 'H') - - Returns - ------ - Class object - """ - - plot_obj = Visualize(element, self.results) - - return plot_obj - def index_to_time(self, index): """ Returns the time value or array given the index/indices @@ -1358,213 +1365,3 @@ def time_to_index(self, time): index = (np.abs(self.results.time.value - time)).argmin() return index - -class Visualize(NEI): - """ - Store plotting results from the simulation - """ - def __init__(self, element, results): - self.element = element - self._results = results - - def index_to_time(self, index): - """ - Inherits the index_to_time method of the NEI class - """ - return super(Visualize, self).index_to_time(index) - - - def ionicfrac_evol_plot(self,x_axis, ion='all'): - """ - Creates a plot of the ionic fraction time evolution of element inputs - - Paramaters - ------ - element: string, - The elemental symbal of the atom (i.e. 'H'), - ion: array-like, dtype=int - The repective integer charge of the atomic particle (i.e. 0 or [0,1]) - - x_axis: ~astropy.units.Quantity: array-like , - The xaxis to plot the ionic fraction evolution over. - Can only be distance or time. - - """ - #Check if element input is a string - if not isinstance(self.element, str): - raise TypeError('The element input must be a string') - - - #Check to see if x_axis units are of time or length - if isinstance(x_axis, u.Quantity): - try: - x = x_axis.to(u.s) - xlabel = 'Time' - except Exception: - x = x_axis.to(u.R_sun) - xlabel = 'Distance' - else: - raise TypeError('Invalid x-axis units. Must be units of length or time.') - - - if ion == 'all': - charge_states = pl.atomic.atomic_number(self.element) + 1 - else: - #Ensure ion input is array of integers - charge_states = np.array(ion, dtype=np.int16) - - linsetyles = ['--','-.',':','-'] - - if ion =='all': - for nstate in range(charge_states): - ionic_frac = self.results.ionic_fractions[self.element][:,nstate] - plt.plot(x.value,ionic_frac, linestyle= linsetyles[nstate % len(linsetyles)], label=f'{self.element}+{nstate}') - plt.xlabel(f'{xlabel} ({x.unit})') - plt.ylabel('Ionic Fraction') - plt.title('Ionic Fraction Evolution of {}'.format(self.element)) - plt.legend(loc='center left') - else: - if charge_states.size > 1: - for nstate in charge_states: - ionic_frac = self.results.ionic_fractions[self.element][:,nstate] - plt.plot(x.value,ionic_frac, linestyle= linsetyles[nstate % len(linsetyles)], label=f'{self.element}+{nstate}') - plt.xlabel(f'{xlabel} ({x.unit})') - plt.ylabel('Ionic Fraction') - plt.title('Ionic Fraction Evolution of {}'.format(self.element)) - plt.legend() - #plt.show() - else: - ionic_frac = self.results.ionic_fractions[self.element][:,charge_states] - plt.plot(x.value,ionic_frac) - plt.xlabel(f'{xlabel} ({x.unit})') - plt.ylabel('Ionic Fraction') - plt.title('Ionic Fraction Evolution of %s$^{%i+}$'%(self.element,ion)) - #plt.show() - - def ionicfrac_bar_plot(self, time_index): - """ - Creates a bar plot of the ion fraction change at a particular time index - - Parameters - ------ - time_index: int, - The particular time index at which to collect the various ion fractiom - change - """ - - if not isinstance(self.element, str): - raise TypeError('The element input must be a string') - - ion = pl.atomic.atomic_number(self.element) - - charge_states = np.linspace(0, ion, ion+1, dtype=np.int16) - - width=1.0 - - fig, ax = plt.subplots() - if isinstance(time_index, (list, np.ndarray)): - - alpha = 1.0 - colors = ['blue', 'red'] - - #Color index counter - color_idx = 1 - - for idx in time_index: - - #Toggle between zero and one for colors array - color_idx ^= 1 - - ax.bar(charge_states, self.results.ionic_fractions[self.element][idx,:], alpha=alpha, \ - width=width, color=colors[color_idx], - label='Time:{time:.{number}f}'.format(time=self.index_to_time(idx), number=1)) - alpha -= 0.2 - ax.set_xticks(charge_states-width/2.0) - ax.set_xticklabels(charge_states) - ax.set_title(f'{self.element}') - - ax.set_xlabel('Charge State') - ax.set_ylabel('Ionic Fraction') - - ax.legend(loc='best') - #plt.show() - - else: - ax.bar(x, self.results.ionic_fractions[self.element][time_index,:], alpha=1.0, width=width) - ax.set_xticks(charge_states-width/2.0) - ax.set_xticklabels(charge_states) - ax.set_title(f'{self.element}') - ax.set_xlabel('Charge State') - ax.set_ylabel('Ionic Fraction') - #plt.show() - - def rh_density_plot(self, gamma, mach, ion='None'): - """ - Creates a plot of the Rankine-Huguniot jump relation for the - density of our element - - Parameters - ------ - gamma: float, - The specific heats ratio of the system - mach: float, - The mach number of the scenario - ion: int, - The ionic integer charge of the element in question - """ - - #Instantiate the MHD class - mhd = shocks.MHD() - - nstates = pl.atomic.atomic_number(self.element) + 1 - - if ion == 'None': - - for charge in range(nstates): - - post_rho = mhd.rh_density(self.results.number_densities[self.element].value[0, charge], gamma, mach) - - plt.semilogy(post_rho, label=f'{self.element}{charge}+') - - plt.legend() - #plt.show() - - else: - - if not isinstance(ion, int): - raise TypeError('Please make sure that your charge value is an integer') - elif ion > nstates: - raise ValueError('The ionic charge input is greater than allowed for this element') - - - post_rho = mhd.rh_density(init_dens = self.results.number_densities[self.element].value[0, ion], gamma=gamma, mach=mach) - - plt.semilogy(post_rho) - plt.title('Density Shock Transition') - plt.xlabel('Mach Number') - plt.ylabel('Number Density') - #plt.show() - - def rh_temp_plot(self, gamma, mach): - """ - Creates a plot of the Rankine-Huguniot jump relation for the - density of our element - - Parameters - ------ - gamma: float, - The specific heats ratio of the system - mach: float, - The mach number of the scenario - """ - - #Instantiate the MHD class - mhd = shocks.MHD() - - post_temp = mhd.rh_temp(self.results.T_e.value[0], gamma, mach) - - plt.semilogy(mach, post_temp) - plt.title('Temperature Shock Transition') - plt.xlabel('Mach Number') - plt.ylabel('Log Temperature (K)') - #plt.show() diff --git a/plasmapy_nei/nei/tests/test_nei.py b/plasmapy_nei/nei/tests/test_nei.py index d5851e0..1f53944 100644 --- a/plasmapy_nei/nei/tests/test_nei.py +++ b/plasmapy_nei/nei/tests/test_nei.py @@ -1,20 +1,14 @@ """Tests for non-equilibrium ionization modeling classes.""" - -def test_import(): - """Minimal test that importing this works.""" - from plasmapy_nei.nei import NEI - - import astropy.units as u -from ..ionization_states import IonizationStates, particle_symbol -from ..nei import NEI -from ..eigenvaluetable import EigenData2 +from plasmapy.atomic import IonizationStates, particle_symbol +from plasmapy_nei.nei import NEI +from plasmapy_nei.eigen import EigenData import numpy as np import pytest -inputs_dict = {'H': [0.9, 0.1], 'He': [0.5, 0.3, 0.2]} -abundances = {'H': 1, 'He': 0.1} +inputs_dict = {"H": [0.9, 0.1], "He": [0.5, 0.3, 0.2]} +abundances = {"H": 1, "He": 0.1} time_array = np.array([0, 800]) * u.s T_e_array = np.array([4e4, 6e4]) * u.K @@ -22,125 +16,124 @@ def test_import(): tests = { - - 'basic': { - 'inputs': inputs_dict, - 'abundances': abundances, - 'T_e': T_e_array, - 'n': n_array, - 'time_input': time_array, - 'time_start': 0 * u.s, - 'time_max': 800 * u.s, - 'max_steps': 1, - 'dt': 800 * u.s, - 'adapt_dt': False, - 'verbose': True, + "basic": { + "inputs": inputs_dict, + "abundances": abundances, + "T_e": T_e_array, + "n": n_array, + "time_input": time_array, + "time_start": 0 * u.s, + "time_max": 800 * u.s, + "max_steps": 1, + "dt": 800 * u.s, + "adapt_dt": False, + "verbose": True, }, - - 'T_e constant': { - 'inputs': inputs_dict, - 'abundances': abundances, - 'T_e': 1 * u.MK, - 'n': n_array, - 'time_input': time_array, - 'time_start': 0 * u.s, - 'time_max': 800 * u.s, - 'dt': 100 * u.s, - 'max_steps': 2, - 'adapt_dt': False, - 'verbose': True, + "T_e constant": { + "inputs": inputs_dict, + "abundances": abundances, + "T_e": 1 * u.MK, + "n": n_array, + "time_input": time_array, + "time_start": 0 * u.s, + "time_max": 800 * u.s, + "dt": 100 * u.s, + "max_steps": 2, + "adapt_dt": False, + "verbose": True, }, - - 'n_e constant': { - 'inputs': inputs_dict, - 'abundances': abundances, - 'T_e': T_e_array, - 'n': 1e9 * u.cm ** -3, - 'time_input': time_array, - 'time_start': 0 * u.s, - 'time_max': 800 * u.s, - 'max_steps': 2, - 'adapt_dt': False, - 'dt': 100 * u.s, - 'verbose': True, + "n_e constant": { + "inputs": inputs_dict, + "abundances": abundances, + "T_e": T_e_array, + "n": 1e9 * u.cm ** -3, + "time_input": time_array, + "time_start": 0 * u.s, + "time_max": 800 * u.s, + "max_steps": 2, + "adapt_dt": False, + "dt": 100 * u.s, + "verbose": True, }, - - 'T_e function': { - 'inputs': inputs_dict, - 'abundances': abundances, - 'T_e': lambda time: 1e4 * (1 + time/u.s) * u.K, - 'n': 1e15 * u.cm **-3, - 'time_max': 800 * u.s, - 'max_steps': 2, - 'dt': 100 * u.s, - 'adapt_dt': False, - 'verbose': True, + "T_e function": { + "inputs": inputs_dict, + "abundances": abundances, + "T_e": lambda time: 1e4 * (1 + time / u.s) * u.K, + "n": 1e15 * u.cm ** -3, + "time_max": 800 * u.s, + "max_steps": 2, + "dt": 100 * u.s, + "adapt_dt": False, + "verbose": True, }, - - 'n function': { - 'inputs': inputs_dict, - 'abundances': abundances, - 'T_e': 6e4 * u.K, - 'n': lambda time: 1e9 * (1 + time/u.s) * u.cm ** -3, - 'time_start': 0 * u.s, - 'time_max': 800 * u.s, - 'adapt_dt': False, - 'dt': 200 * u.s, - 'verbose': True, - 'max_steps': 4 + "n function": { + "inputs": inputs_dict, + "abundances": abundances, + "T_e": 6e4 * u.K, + "n": lambda time: 1e9 * (1 + time / u.s) * u.cm ** -3, + "time_start": 0 * u.s, + "time_max": 800 * u.s, + "adapt_dt": False, + "dt": 200 * u.s, + "verbose": True, + "max_steps": 4, }, - - 'equil test cool': { - 'inputs': ['H', 'He', 'N'], - 'abundances': {'H': 1, 'He': 0.1, 'C': 1e-4, 'N': 1e-4, 'O': 1e-4, 'Fe': 1e-4}, - 'T_e': 10001.0 * u.K, - 'n': 1e13 * u.cm ** -3, - 'time_max': 2e6 * u.s, - 'tol': 1e-9, - 'adapt_dt': False, - 'dt': 5e5 * u.s, - 'max_steps': 4, - 'verbose': True, + "equil test cool": { + "inputs": ["H", "He", "N"], + "abundances": {"H": 1, "He": 0.1, "C": 1e-4, "N": 1e-4, "O": 1e-4, "Fe": 1e-4}, + "T_e": 10001.0 * u.K, + "n": 1e13 * u.cm ** -3, + "time_max": 2e6 * u.s, + "tol": 1e-9, + "adapt_dt": False, + "dt": 5e5 * u.s, + "max_steps": 4, + "verbose": True, }, - - 'equil test hot': { - 'inputs': ['H', 'He', 'C'], - 'abundances': {'H': 1, 'He': 0.1, 'C': 1e-4, 'N': 1e-4, 'O': 1e-4, 'Fe': 1e-4, 'S': 2e-6}, - 'T_e': 7e6 * u.K, - 'n': 1e9 * u.cm ** -3, - 'time_max': 1e8 * u.s, - 'dt': 5e7 * u.s, - 'max_steps': 3, - 'adapt_dt': False, - 'verbose': True, + "equil test hot": { + "inputs": ["H", "He", "C"], + "abundances": { + "H": 1, + "He": 0.1, + "C": 1e-4, + "N": 1e-4, + "O": 1e-4, + "Fe": 1e-4, + "S": 2e-6, + }, + "T_e": 7e6 * u.K, + "n": 1e9 * u.cm ** -3, + "time_max": 1e8 * u.s, + "dt": 5e7 * u.s, + "max_steps": 3, + "adapt_dt": False, + "verbose": True, }, - - 'equil test start far out of equil': { - 'inputs': { - 'H': [0.99, 0.01], - 'He': [0.5, 0.0, 0.5], - 'O': [0.2, 0, 0.2, 0, 0.2, 0, 0.2, 0, 0.2], + "equil test start far out of equil": { + "inputs": { + "H": [0.99, 0.01], + "He": [0.5, 0.0, 0.5], + "O": [0.2, 0, 0.2, 0, 0.2, 0, 0.2, 0, 0.2], }, - 'abundances': {'H': 1, 'He': 0.1, 'O': 1e-4}, - 'T_e': 3e6 * u.K, - 'n': 1e9 * u.cm ** -3, - 'dt': 1e6 * u.s, - 'time_start': 0 * u.s, - 'time_max': 1e6 * u.s, - 'adapt_dt': False, - 'verbose': True, - 'max_steps': 2 + "abundances": {"H": 1, "He": 0.1, "O": 1e-4}, + "T_e": 3e6 * u.K, + "n": 1e9 * u.cm ** -3, + "dt": 1e6 * u.s, + "time_start": 0 * u.s, + "time_max": 1e6 * u.s, + "adapt_dt": False, + "verbose": True, + "max_steps": 2, }, - - 'adapt dt': { - 'inputs': ['H', 'He'], - 'abundances': {'H': 1, 'He': 0.1}, - 'T_e': lambda t: u.K * (1e6 + 1.3e4*np.sin(t.value)), - 'n': 1e10 * u.cm ** -3, - 'max_steps': 300, - 'time_start': 0 * u.s, - 'time_max': 2*np.pi * u.s, - 'adapt_dt': True, + "adapt dt": { + "inputs": ["H", "He"], + "abundances": {"H": 1, "He": 0.1}, + "T_e": lambda t: u.K * (1e6 + 1.3e4 * np.sin(t.value)), + "n": 1e10 * u.cm ** -3, + "max_steps": 300, + "time_start": 0 * u.s, + "time_max": 2 * np.pi * u.s, + "adapt_dt": True, }, } @@ -148,12 +141,11 @@ def test_import(): class TestNEI: - @classmethod def setup_class(cls): cls.instances = {} - @pytest.mark.parametrize('test_name', test_names) + @pytest.mark.parametrize("test_name", test_names) def test_instantiate(self, test_name): try: instance = NEI(**tests[test_name]) @@ -161,46 +153,46 @@ def test_instantiate(self, test_name): except Exception as exc: raise Exception(f"Problem with test {test_name}") from exc - @pytest.mark.parametrize('test_name', test_names) + @pytest.mark.parametrize("test_name", test_names) def test_time_start(self, test_name): instance = self.instances[test_name] - if 'time_start' in tests[test_name].keys(): - assert tests[test_name]['time_start'] == instance.time_start - elif 'time_input' in tests[test_name].keys(): - assert tests[test_name]['time_input'].min() == instance.time_start + if "time_start" in tests[test_name].keys(): + assert tests[test_name]["time_start"] == instance.time_start + elif "time_input" in tests[test_name].keys(): + assert tests[test_name]["time_input"].min() == instance.time_start else: assert instance.time_start == 0 * u.s - @pytest.mark.parametrize('test_name', test_names) + @pytest.mark.parametrize("test_name", test_names) def test_time_max(self, test_name): instance = self.instances[test_name] - if 'time_max' in tests[test_name].keys(): - assert tests[test_name]['time_max'] == instance.time_max + if "time_max" in tests[test_name].keys(): + assert tests[test_name]["time_max"] == instance.time_max - @pytest.mark.parametrize('test_name', test_names) + @pytest.mark.parametrize("test_name", test_names) def test_initial_type(self, test_name): instance = self.instances[test_name] assert isinstance(instance.initial, IonizationStates) - @pytest.mark.parametrize('test_name', test_names) + @pytest.mark.parametrize("test_name", test_names) def test_n_input(self, test_name): actual = self.instances[test_name].n_input - expected = tests[test_name]['n'] + expected = tests[test_name]["n"] if isinstance(expected, u.Quantity) and not expected.isscalar: assert all(expected == actual) else: assert expected == actual - @pytest.mark.parametrize('test_name', test_names) + @pytest.mark.parametrize("test_name", test_names) def test_T_e_input(self, test_name): actual = self.instances[test_name].T_e_input - expected = tests[test_name]['T_e'] + expected = tests[test_name]["T_e"] if isinstance(expected, u.Quantity) and not expected.isscalar: assert all(expected == actual) else: assert expected == actual - @pytest.mark.parametrize('test_name', test_names) + @pytest.mark.parametrize("test_name", test_names) def test_electron_temperature(self, test_name): instance = self.instances[test_name] T_e_input = instance.T_e_input @@ -217,31 +209,38 @@ def test_electron_temperature(self, test_name): assert np.isclose(T_e.value, T_e_func.value) if callable(T_e_input): - assert instance.T_e_input(instance.time_start) == \ - instance.electron_temperature(instance.time_start) + assert instance.T_e_input( + instance.time_start + ) == instance.electron_temperature(instance.time_start) @pytest.mark.parametrize( - 'test_name', - [test_name for test_name in test_names if isinstance(tests[test_name]['inputs'], dict)], + "test_name", + [ + test_name + for test_name in test_names + if isinstance(tests[test_name]["inputs"], dict) + ], ) def test_initial_ionfracs(self, test_name): - original_inputs = tests[test_name]['inputs'] + original_inputs = tests[test_name]["inputs"] original_elements = original_inputs.keys() for element in original_elements: assert np.allclose( original_inputs[element], - self.instances[test_name].initial.ionic_fractions[particle_symbol(element)] + self.instances[test_name].initial.ionic_fractions[ + particle_symbol(element) + ], ) - @pytest.mark.parametrize('test_name', test_names) + @pytest.mark.parametrize("test_name", test_names) def test_simulate(self, test_name): try: self.instances[test_name].simulate() except Exception as exc: raise ValueError(f"Unable to simulate for test: {test_name}") from exc - @pytest.mark.parametrize('test_name', test_names) + @pytest.mark.parametrize("test_name", test_names) def test_simulation_end(self, test_name): instance = self.instances[test_name] time = instance.results.time @@ -249,7 +248,7 @@ def test_simulation_end(self, test_name): max_steps = instance.max_steps if np.isnan(end_time.value): - raise Exception('End time is NaN.') + raise Exception("End time is NaN.") got_to_max_steps = len(time) == instance.max_steps + 1 got_to_time_max = np.isclose(time[-1].value, instance.time_max.value) @@ -260,12 +259,11 @@ def test_simulation_end(self, test_name): if not got_to_max_steps and not got_to_time_max: print(f"time = {time}") print(f"max_steps = {max_steps}") - print(f'time_max = {instance.time_max}') + print(f"time_max = {instance.time_max}") raise Exception("Problem with end time.") @pytest.mark.parametrize( - 'test_name', - [test_name for test_name in test_names if 'equil' in test_name] + "test_name", [test_name for test_name in test_names if "equil" in test_name] ) def test_equilibration(self, test_name): """ @@ -273,31 +271,36 @@ def test_equilibration(self, test_name): """ instance = self.instances[test_name] T_e = instance.T_e_input - assert isinstance(T_e, u.Quantity) and T_e.isscalar, \ - "This test can only be used for cases where T_e is constant." + assert ( + isinstance(T_e, u.Quantity) and T_e.isscalar + ), "This test can only be used for cases where T_e is constant." equil_dict = instance.equil_ionic_fractions(T_e) for element in instance.elements: assert np.allclose( equil_dict[element], instance.results.ionic_fractions[element][-1, :] ) - @pytest.mark.parametrize('test_name', test_names) + @pytest.mark.parametrize("test_name", test_names) def test_initial_results(self, test_name): initial = self.instances[test_name].initial results = self.instances[test_name].results - assert initial.elements == results.elements + # Make sure that the elements are equal to each other + assert initial.ionic_fractions.keys() == results.ionic_fractions.keys() assert initial.abundances == results.abundances - for elem in initial.elements: - assert np.allclose(results.ionic_fractions[elem][0, :], initial.ionic_fractions[elem]) + for elem in initial.ionic_fractions.keys(): # TODO: enable initial.elements + assert np.allclose( + results.ionic_fractions[elem][0, :], initial.ionic_fractions[elem] + ) - @pytest.mark.parametrize('test_name', test_names) + @pytest.mark.parametrize("test_name", test_names) def test_final_results(self, test_name): - #initial = self.instances[test_name].initial + # initial = self.instances[test_name].initial final = self.instances[test_name].final results = self.instances[test_name].results - assert final.elements == results.elements + # Make sure the elements are equal to each other + assert final.ionic_fractions.keys() == results.ionic_fractions.keys() assert final.abundances == results.abundances - for elem in final.elements: + for elem in final.ionic_fractions.keys(): assert np.allclose( results.ionic_fractions[elem][-1, :], final.ionic_fractions[elem] ) From 13eb41d6bbf573bdc7c9be6d79e29f213c9258bf Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Fri, 1 May 2020 11:09:29 -0400 Subject: [PATCH 13/49] Use black to clean up code --- plasmapy_nei/nei/nei.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/plasmapy_nei/nei/nei.py b/plasmapy_nei/nei/nei.py index 43bb4ea..2fb9903 100644 --- a/plasmapy_nei/nei/nei.py +++ b/plasmapy_nei/nei/nei.py @@ -13,7 +13,6 @@ import warnings - # TODO: Allow this to keep track of velocity and position too, and # eventually to have density and temperature be able to be functions of # position. (and more complicated expressions for density and @@ -34,6 +33,7 @@ # to put in an `elements` attribute in IonizationStates, and # should be corrected. + class NEIError(Exception): pass @@ -479,11 +479,7 @@ def __init__( n_init = self.hydrogen_number_density(self.time_start) self.initial = IonizationStates( - inputs=inputs, - abundances=abundances, - T_e=T_e_init, - n=n_init, - tol=tol, + inputs=inputs, abundances=abundances, T_e=T_e_init, n=n_init, tol=tol, ) self.tol = tol @@ -988,7 +984,9 @@ def initial(self) -> IonizationStates: def initial(self, initial_states: Optional[IonizationStates]): if isinstance(initial_states, IonizationStates): self._initial = initial_states - self._elements = initial_states.ionic_fractions.keys() # TODO IonizationStates + self._elements = ( + initial_states.ionic_fractions.keys() + ) # TODO IonizationStates elif initial_states is None: self._ionstates = None else: @@ -1071,9 +1069,7 @@ def _finalize_simulation(self): self._final = IonizationStates( inputs=final_ionfracs, abundances=self.abundances, - n=np.sum( - self.results.number_densities["H"][-1, :] - ), # modify this later?, + n=np.sum(self.results.number_densities["H"][-1, :]), # modify this later?, T_e=self.results.T_e[-1], tol=1e-6, ) From eae0d555e9d086386c89ce660ddff51a1b5dad2f Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Fri, 1 May 2020 11:29:39 -0400 Subject: [PATCH 14/49] Rename Simulation, add docstrings, etc. --- plasmapy_nei/__init__.py | 1 + plasmapy_nei/nei/__init__.py | 2 +- plasmapy_nei/nei/nei.py | 18 ++++++++---------- 3 files changed, 10 insertions(+), 11 deletions(-) diff --git a/plasmapy_nei/__init__.py b/plasmapy_nei/__init__.py index 4c37b70..4670952 100644 --- a/plasmapy_nei/__init__.py +++ b/plasmapy_nei/__init__.py @@ -4,6 +4,7 @@ try: from .version import __version__ + del version except Exception as exc: warnings.warn("Unable to import __version__") diff --git a/plasmapy_nei/nei/__init__.py b/plasmapy_nei/nei/__init__.py index 239364c..8928a9a 100644 --- a/plasmapy_nei/nei/__init__.py +++ b/plasmapy_nei/nei/__init__.py @@ -1,3 +1,3 @@ """Classes that perform non-equilibrium ionization modeling""" -from plasmapy_nei.nei.nei import NEI, NEIError +from plasmapy_nei.nei.nei import NEI, NEIError, SimulationResults diff --git a/plasmapy_nei/nei/nei.py b/plasmapy_nei/nei/nei.py index 2fb9903..89a65fe 100644 --- a/plasmapy_nei/nei/nei.py +++ b/plasmapy_nei/nei/nei.py @@ -1,6 +1,6 @@ """Contains classes to represent non-equilibrium ionization simulations.""" -__all__ = ["NEI", "NEIError"] +__all__ = ["NEI", "NEIError", "SimulationResults"] import numpy as np @@ -38,7 +38,7 @@ class NEIError(Exception): pass -class Simulation: +class SimulationResults: """ Results from a non-equilibrium ionization simulation. @@ -245,7 +245,6 @@ def abundances(self) -> Dict[str, float]: The keys are the atomic symbols and the values are a `float` representing that element's elemental abundance. - """ return self._abundances @@ -257,7 +256,6 @@ def ionic_fractions(self) -> Dict[str, np.ndarray]: The keys of this dictionary are atomic symbols. The values are 2D arrays where the first index refers to the time step and the second index refers to the integer charge. - """ return self._ionic_fractions @@ -430,9 +428,8 @@ class NEI: Notes ----- The ionization and recombination rates are from Chianti version - 8.x. These rates include radiative and dielectronic recombination. + 8.7. These rates include radiative and dielectronic recombination. Photoionization is not included. - """ def __init__( @@ -552,7 +549,6 @@ def equil_ionic_fractions( neither `T_e` or `time` is provided and the temperature for the simulation is given by a constant, the this method will assume that `T_e` is the temperature of the simulation. - """ if T_e is not None and time is not None: @@ -586,6 +582,7 @@ def equil_ionic_fractions( @property def elements(self) -> List[str]: + """A `list` of the elements.""" return self._elements @elements.setter @@ -741,6 +738,7 @@ def dt_input(self, dt: Optional[u.Quantity]): @property def dt_min(self) -> u.Quantity: + """The minimum time step.""" return self._dt_min @dt_min.setter @@ -993,7 +991,7 @@ def initial(self, initial_states: Optional[IonizationStates]): raise TypeError("Expecting an IonizationStates instance.") @property - def results(self) -> Simulation: + def results(self) -> SimulationResults: """ Return the `~nei.classes.Simulation` class instance that corresponds to the simulation results. @@ -1017,7 +1015,7 @@ def final(self) -> IonizationStates: def _initialize_simulation(self): - self._results = Simulation( + self._results = SimulationResults( initial=self.initial, n_init=self.hydrogen_number_density(self.time_start), T_e_init=self.electron_temperature(self.time_start), @@ -1027,7 +1025,7 @@ def _initialize_simulation(self): self._old_time = self.time_start.to(u.s) self._new_time = self.time_start.to(u.s) - def simulate(self) -> Simulation: + def simulate(self) -> SimulationResults: """ Perform a non-equilibrium ionization simulation. From b1dd5c75bc81361f7715fd4b381f0de2c90cf79d Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Fri, 1 May 2020 13:40:26 -0400 Subject: [PATCH 15/49] Create minimal docs for nei & eigen subpackages --- docs/eigen/index.rst | 10 ++++++++++ docs/index.rst | 3 +++ docs/nei/index.rst | 10 ++++++++++ 3 files changed, 23 insertions(+) create mode 100644 docs/eigen/index.rst create mode 100644 docs/nei/index.rst diff --git a/docs/eigen/index.rst b/docs/eigen/index.rst new file mode 100644 index 0000000..bc683fc --- /dev/null +++ b/docs/eigen/index.rst @@ -0,0 +1,10 @@ +.. _eigen: + +********************** +``plasmapy_nei.eigen`` +********************** + +.. py:currentmodule:: plasmapy_nei.eigen + +.. automodapi:: plasmapy_nei.eigen + :no-heading: diff --git a/docs/index.rst b/docs/index.rst index b96b176..c76c281 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -28,6 +28,9 @@ The early versions of this package will account for collisional ionization, radi :maxdepth: 2 :caption: Contents: + ``plasmapy_nei.eigen`` + ``plasmapy_nei.nei`` + Indices and tables ================== diff --git a/docs/nei/index.rst b/docs/nei/index.rst new file mode 100644 index 0000000..7a9c707 --- /dev/null +++ b/docs/nei/index.rst @@ -0,0 +1,10 @@ +.. _nei: + +******************** +``plasmapy_nei.nei`` +******************** + +.. py:currentmodule:: plasmapy_nei.nei + +.. automodapi:: plasmapy_nei.nei + :no-heading: From 7511e41b1aa2a7521da4a751f7025645604b9c04 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Fri, 1 May 2020 14:21:38 -0400 Subject: [PATCH 16/49] Update ReST formatting --- plasmapy_nei/nei/nei.py | 143 ++++++++++++++++++++-------------------- 1 file changed, 71 insertions(+), 72 deletions(-) diff --git a/plasmapy_nei/nei/nei.py b/plasmapy_nei/nei/nei.py index 89a65fe..589107c 100644 --- a/plasmapy_nei/nei/nei.py +++ b/plasmapy_nei/nei/nei.py @@ -122,25 +122,25 @@ def _assign( Parameters ---------- - new_time: ~astropy.units.Quantity + new_time The time associated with this time step. new_ionfracs: dict The new ionization fractions for this time step. The keys of this `dict` are the atomic symbols of the elements being tracked, and with the corresponding value being an - `~numpy.ndarray` representing the ionic fractions. Each + ``numpy.ndarray`` representing the ionic fractions. Each element's array must have a length of the atomic number plus one, and be normalized to one with all values between zero and one. - new_n: ~astropy.units.Quantity + new_n The new number density scaling factor for this time step. The number densities of each ionic species will be the product of this scaling factor, the element's abundance, and - the ionic fraction given in `new_ionfracs`. + the ionic fraction given in ``new_ionfracs``. - new_T_e: ~astropy.units.Quantity + new_T_e The new electron temperature. """ @@ -190,7 +190,7 @@ def _cleanup(self): This method removes the excess elements from each array that did not end up getting used for a time step in the simulation - and sets the `last_step` attribute. + and sets the ``last_step`` attribute. """ nsteps = self._index @@ -264,7 +264,7 @@ def number_densities(self) -> Dict[str, u.Quantity]: """ Return the number densities over the course of the simulation. - The keys of `number_densities` are atomic symbols. The values + The keys of ``number_densities`` are atomic symbols. The values are 2D arrays with units of number density where the first index refers to the time step and the second index is the integer charge. @@ -278,7 +278,7 @@ def n_elem(self) -> Dict[str, u.Quantity]: The number densities of each element over the course of the simulation. - The keys of `n_elem` are atomic symbols. The values are 1D + The keys of ``n_elem`` are atomic symbols. The values are 1D arrays with units of number density where the index refers to the time step. @@ -322,48 +322,48 @@ class NEI: ---------- inputs - T_e: `~astropy.units.Quantity` or `callable` + T_e: astropy.units.Quantity or `callable` The electron temperature, which may be a constant, an array of temperatures corresponding to the times in `time_input`, or a function that yields the temperature as a function of time. - n: `~astropy.units.Quantity` or `callable` + n: astropy.units.Quantity or `callable` The number density multiplicative factor. The number density of - each element will be `n` times the abundance given in - `abundances`. For example, if `abundance['H'] = 1`, then this + each element will be ``n`` times the abundance given in + ``abundances``. For example, if ``abundance['H'] = 1``, then this will correspond to the number density of hydrogen (including neutral hydrogen and protons). This factor may be a constant, an array of number densities over time, or a function that yields a number density as a function of time. - time_input: ~astropy.units.Quantity, optional - An array containing the times associated with `n` and `T_e` in + time_input: astropy.units.Quantity, optional + An array containing the times associated with ``n`` and ``T_e`` in units of time. - time_start: ~astropy.units.Quantity, optional + time_start: astropy.units.Quantity, optional The start time for the simulation. If density and/or temperature are given by arrays, then this argument must be - greater than `time_input[0]`. If this argument is not supplied, - then `time_start` defaults to `time_input[0]` (if given) and + greater than ``time_input[0]``. If this argument is not supplied, + then ``time_start`` defaults to ``time_input[0]`` (if given) and zero seconds otherwise. - time_max: ~astropy.units.Quantity + time_max: astropy.units.Quantity The maximum time for the simulation. If density and/or temperature are given by arrays, then this argument must be less - than `time_input[-1]`. + than ``time_input[-1]``. max_steps: `int` The maximum number of time steps to be taken during a simulation. - dt: `~astropy.units.Quantity` - The time step. If `adapt_dt` is `False`, then `dt` is the time - step for the whole simulation. + dt: astropy.units.Quantity + The time step. If ``adapt_dt`` is `False`, then ``dt`` is the + time step for the whole simulation. - dt_max: `~astropy.units.Quantity` + dt_max: astropy.units.Quantity The maximum time step to be used with an adaptive time step. - dt_min: `~astropy.units.Quantity` + dt_min: astropy.units.Quantity The minimum time step to be used with an adaptive time step. adapt_dt: `bool` @@ -373,7 +373,7 @@ class NEI: safety_factor: `float` or `int` A multiplicative factor to multiply by the time step when - `adapt_dt` is `True`. Lower values improve accuracy, whereas + ``adapt_dt`` is `True`. Lower values improve accuracy, whereas higher values reduce computational time. Not yet implemented. tol: float @@ -381,7 +381,7 @@ class NEI: verbose: bool, optional A flag stating whether or not to print out information for every - time step. Setting `verbose` to `True` is useful for testing. + time step. Setting ``verbose`` to `True` is useful for testing. Defaults to `False`. abundances: dict @@ -408,12 +408,12 @@ class NEI: >>> results = sim.simulate() - The initial results are stored in the `initial` attribute. + The initial results are stored in the ``initial`` attribute. >>> sim.initial.ionic_fractions['H'] array([0.9, 0.1]) - The final results can be access with the `final` attribute. + The final results can be access with the ``final`` attribute. >>> sim.final.ionic_fractions['H'] array([0.16665179, 0.83334821]) @@ -422,7 +422,7 @@ class NEI: >>> sim.final.T_e - Both `initial` and `final` are instances of the `IonizationStates` + Both ``initial`` and ``final`` are instances of the ``IonizationStates`` class. Notes @@ -530,25 +530,25 @@ def equil_ionic_fractions( Parameters ---------- - T_e: ~astropy.units.Quantity, optional + T_e: astropy.units.Quantity, optional The electron temperature in units that can be converted to kelvin. - time: ~astropy.units.Quantity, optional + time: astropy.units.Quantity, optional The time in units that can be converted to seconds. Returns ------- - equil_ionfracs: dict + equil_ionfracs: `dict` The equilibrium ionic fractions for the elements contained within this class Notes ----- - Only one of `T_e` and `time` may be included as an argument. If - neither `T_e` or `time` is provided and the temperature for the - simulation is given by a constant, the this method will assume - that `T_e` is the temperature of the simulation. + Only one of ``T_e`` and ``time`` may be included as an argument. + If neither ``T_e`` or ``time`` is provided and the temperature + for the simulation is given by a constant, the this method will + assume that ``T_e`` is the temperature of the simulation. """ if T_e is not None and time is not None: @@ -619,17 +619,17 @@ def tol(self, value: Union[float, int]): try: value = float(value) except Exception as exc: - raise TypeError("Invalid tolerance.") + raise TypeError(f"Invalid tolerance: {value}") from exc if not 0 <= value < 1: raise ValueError("Need 0 <= tol < 1.") self._tol = value @property - def time_input(self) -> Optional[u.Quantity]: + def time_input(self) -> Optional[u.s]: return self._time_input @time_input.setter - def time_input(self, times: Optional[u.Quantity]): + def time_input(self, times: Optional[u.s]): if times is None: self._time_input = None elif isinstance(times, u.Quantity): @@ -646,12 +646,12 @@ def time_input(self, times: Optional[u.Quantity]): raise TypeError("Invalid time_input.") @property - def time_start(self) -> u.Quantity: + def time_start(self) -> u.s: """The start time of the simulation.""" return self._time_start @time_start.setter - def time_start(self, time: u.Quantity): + def time_start(self, time: u.s): if time is None: self._time_start = 0.0 * u.s elif isinstance(time, u.Quantity): @@ -674,12 +674,12 @@ def time_start(self, time: u.Quantity): raise TypeError("Invalid time_start.") from None @property - def time_max(self) -> u.Quantity: + def time_max(self) -> u.s: """The maximum time allowed for the simulation.""" return self._time_max @time_max.setter - def time_max(self, time: u.Quantity): + def time_max(self, time: u.s): if time is None: self._time_max = ( self.time_input[-1] if self.time_input is not None else np.inf * u.s @@ -720,12 +720,12 @@ def adapt_dt(self, choice: Optional[bool]): raise TypeError("Invalid value for adapt_dt") @property - def dt_input(self) -> Optional[u.Quantity]: + def dt_input(self) -> Optional[u.s]: """Return the inputted time step.""" return self._dt_input @dt_input.setter - def dt_input(self, dt: Optional[u.Quantity]): + def dt_input(self, dt: Optional[u.s]): if dt is None: self._dt_input = None elif isinstance(dt, u.Quantity): @@ -737,12 +737,12 @@ def dt_input(self, dt: Optional[u.Quantity]): raise NEIError("Invalid dt.") @property - def dt_min(self) -> u.Quantity: + def dt_min(self) -> u.s: """The minimum time step.""" return self._dt_min @dt_min.setter - def dt_min(self, value: u.Quantity): + def dt_min(self, value: u.s): if not isinstance(value, u.Quantity): raise TypeError("dt_min must be a Quantity.") try: @@ -816,16 +816,16 @@ def verbose(self, choice: bool): @u.quantity_input def in_time_interval(self, time: u.s, buffer: u.s = 1e-9 * u.s): """ - Return `True` if the `time` is between `time_start - buffer` and - `time_max + buffer` , and `False` otherwise. + Return `True` if the ``time`` is between ``time_start - buffer`` + and ``time_max + buffer`` , and `False` otherwise. Raises ------ TypeError - If `time` or `buffer` is not a `~astropy.units.Quantity` + If ``time`` or ``buffer`` is not a ``astropy.units.Quantity`` - `~astropy.units.UnitsError` - If `time` or `buffer` is not in units of time. + astropy.units.UnitsError + If ``time`` or ``buffer`` is not in units of time. """ return self.time_start - buffer <= time <= self.time_max + buffer @@ -966,7 +966,8 @@ def hydrogen_number_density(self, time: u.Quantity) -> u.Quantity: @property def EigenDataDict(self) -> Dict[str, EigenData]: """ - Return a `dict` containing `~nei.class + Return a `dict` containing `~plasmapy.eigen.EigenData` instances + for each element. """ return self._EigenDataDict @@ -993,7 +994,7 @@ def initial(self, initial_states: Optional[IonizationStates]): @property def results(self) -> SimulationResults: """ - Return the `~nei.classes.Simulation` class instance that + Return the `~plasmapy_nei.nei.SimulationResults` class instance that corresponds to the simulation results. """ @@ -1031,10 +1032,10 @@ def simulate(self) -> SimulationResults: Returns ------- - results: ~nei.classes.Simulation + results: `~plasmapy_nei.classes.Simulation` The results from the simulation (which are also stored in - the `results` attribute of the `~nei.classes.NEI` instance - this method was called from. + the ``results`` attribute of the `~plasmapy_nei.nei.NEI` + instance this method was called from. """ @@ -1208,25 +1209,25 @@ def set_timestep(self, dt: u.Quantity = None): Notes ----- - If `dt` is not `None`, then the time step will be set to `dt`. + If ``dt`` is not `None`, then the time step will be set to ``dt``. - If `dt` is not set and the `adapt_dt` attribute of an - `~nei.classes.NEI` instance is `True`, then this method will + If ``dt`` is not set and the ``adapt_dt`` attribute of an + `~plasmapy_nei.nei.NEI` instance is `True`, then this method will calculate the time step corresponding to how long it will be until the temperature rises or drops into the next temperature - bin. If this time step is between `dtmin` and `dtmax`, then + bin. If this time step is between ``dtmin`` and ``dtmax``, then - If `dt` is not set and the `adapt_dt` attribute is `False`, then - this method will set the time step as what was inputted to the - `~nei.classes.NEI` class upon instantiation in the `dt` argument - or through the `~nei.classes.NEI` class's `dt_input` attribute. + If ``dt`` is not set and the ``adapt_dt`` attribute is `False`, + then this method will set the time step as what was inputted to + the `~plasmapy_nei.nei.NEI` class upon instantiation in the + ``dt`` argument or through the `~plasmapy_nei.nei.NEI` class's + ``dt_input`` attribute. Raises ------ - ~nei.classes.NEIError - If the time step cannot be set, for example if the `dt` + ~plasmapy_nei.nei.NEIError + If the time step cannot be set, for example if the ``dt`` argument is invalid or the time step cannot be adapted. - """ if dt is not None: @@ -1255,9 +1256,7 @@ def set_timestep(self, dt: u.Quantity = None): self._dt = self._new_time - self._old_time def time_advance(self): - """ - Advance the simulation by one time step. - """ + """Advance the simulation by one time step.""" # TODO: Expand docstring and include equations! # TODO: Fully implement units into this. @@ -1318,7 +1317,7 @@ def time_advance(self): def save(self, filename: str = "nei.h5"): """ - Save the `~nei.classes.NEI` instance to an HDF5 file. Not + Save the `~plasmapy_nei.nei.NEI` instance to an HDF5 file. Not implemented. """ raise NotImplementedError From f4ffc66dbefc7743353bd8c1cbdd8889c6ce8953 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Fri, 1 May 2020 14:27:26 -0400 Subject: [PATCH 17/49] Update type hint annotations Avoid using Optional[u.s] as this seems to cause problems. --- plasmapy_nei/nei/nei.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/plasmapy_nei/nei/nei.py b/plasmapy_nei/nei/nei.py index 589107c..44a4a4b 100644 --- a/plasmapy_nei/nei/nei.py +++ b/plasmapy_nei/nei/nei.py @@ -625,11 +625,11 @@ def tol(self, value: Union[float, int]): self._tol = value @property - def time_input(self) -> Optional[u.s]: + def time_input(self) -> u.s: return self._time_input @time_input.setter - def time_input(self, times: Optional[u.s]): + def time_input(self, times: u.s): if times is None: self._time_input = None elif isinstance(times, u.Quantity): @@ -720,12 +720,12 @@ def adapt_dt(self, choice: Optional[bool]): raise TypeError("Invalid value for adapt_dt") @property - def dt_input(self) -> Optional[u.s]: + def dt_input(self) -> u.s: """Return the inputted time step.""" return self._dt_input @dt_input.setter - def dt_input(self, dt: Optional[u.s]): + def dt_input(self, dt: u.s): if dt is None: self._dt_input = None elif isinstance(dt, u.Quantity): @@ -760,17 +760,17 @@ def dt_min(self, value: u.s): self._dt_min = value @property - def dt_max(self) -> u.Quantity: + def dt_max(self) -> u.s: return self._dt_max @dt_max.setter - def dt_max(self, value: u.Quantity): + def dt_max(self, value: u.s): if not isinstance(value, u.Quantity): raise TypeError("dt_max must be a Quantity.") try: value = value.to(u.s) except u.UnitConversionError as exc: - raise u.UnitConversionError("Invalid units for dt_max.") + raise u.UnitConversionError("Invalid units for dt_max.") from exc if ( hasattr(self, "_dt_input") and self.dt_input is not None @@ -980,7 +980,7 @@ def initial(self) -> IonizationStates: return self._initial @initial.setter - def initial(self, initial_states: Optional[IonizationStates]): + def initial(self, initial_states: IonizationStates): if isinstance(initial_states, IonizationStates): self._initial = initial_states self._elements = ( @@ -1204,7 +1204,7 @@ def set_timestep(self, dt: u.Quantity = None): Parameters ---------- - dt: ~astropy.units.Quantity, optional + dt: astropy.units.Quantity, optional The time step to be used for the next time advance. Notes From 0dd70a7335961f0a8b71b83c21fa5d302ba32295 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Fri, 1 May 2020 14:37:51 -0400 Subject: [PATCH 18/49] Temporarily stop building docstrings from nei subpackage --- docs/nei/index.rst | 3 --- 1 file changed, 3 deletions(-) diff --git a/docs/nei/index.rst b/docs/nei/index.rst index 7a9c707..538412a 100644 --- a/docs/nei/index.rst +++ b/docs/nei/index.rst @@ -5,6 +5,3 @@ ******************** .. py:currentmodule:: plasmapy_nei.nei - -.. automodapi:: plasmapy_nei.nei - :no-heading: From 6f89b13d4fe71bd0031a5d0b031d7587e3b81317 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Fri, 1 May 2020 14:45:05 -0400 Subject: [PATCH 19/49] Update ReST formatting This is to try to pin down some doc build errors for which error messages are not being useful. I largely removed references to external packages (e.g., changing `~astropy.units.Quantity` to ``astropy.units.Quantity``). These should be added back in later. --- docs/nei/index.rst | 9 ++++++--- plasmapy_nei/nei/nei.py | 18 ++++++++---------- 2 files changed, 14 insertions(+), 13 deletions(-) diff --git a/docs/nei/index.rst b/docs/nei/index.rst index 538412a..4558e85 100644 --- a/docs/nei/index.rst +++ b/docs/nei/index.rst @@ -1,7 +1,10 @@ .. _nei: -******************** -``plasmapy_nei.nei`` -******************** +**************** +``plasmapy.nei`` +**************** .. py:currentmodule:: plasmapy_nei.nei + +.. automodapi:: plasmapy_nei.nei + :no-heading: diff --git a/plasmapy_nei/nei/nei.py b/plasmapy_nei/nei/nei.py index 44a4a4b..83c2988 100644 --- a/plasmapy_nei/nei/nei.py +++ b/plasmapy_nei/nei/nei.py @@ -45,21 +45,21 @@ class SimulationResults: Parameters ---------- initial: plasmapy.atomic.IonizationStates - The `IonizationStates` instance representing the ionization + The ``IonizationStates`` instance representing the ionization states of different elements and plasma properties as the initial conditions. - n_init: ~astropy.units.quantity + n_init: astropy.units.quantity The initial number density scaling factor. - T_e_init: ~astropy.units.quantity + T_e_init: astropy.units.quantity The initial electron temperature. max_steps: int The maximum number of time steps that the simulation can take before stopping. - time_start: ~astropy.units.Quantity + time_start: astropy.units.Quantity The time at the start of the simulation. """ @@ -218,9 +218,7 @@ def max_steps(self) -> int: @property def last_step(self) -> int: - """ - The time index of the last step. - """ + """The time index of the last step.""" return self._last_step @property @@ -322,12 +320,12 @@ class NEI: ---------- inputs - T_e: astropy.units.Quantity or `callable` + T_e: astropy.units.Quantity or callable The electron temperature, which may be a constant, an array of temperatures corresponding to the times in `time_input`, or a function that yields the temperature as a function of time. - n: astropy.units.Quantity or `callable` + n: astropy.units.Quantity or callable The number density multiplicative factor. The number density of each element will be ``n`` times the abundance given in ``abundances``. For example, if ``abundance['H'] = 1``, then this @@ -966,7 +964,7 @@ def hydrogen_number_density(self, time: u.Quantity) -> u.Quantity: @property def EigenDataDict(self) -> Dict[str, EigenData]: """ - Return a `dict` containing `~plasmapy.eigen.EigenData` instances + Return a `dict` containing `~plasmapy_nei.eigen.EigenData` instances for each element. """ return self._EigenDataDict From 693603efcdbd154be7e1d9dcd0090042ba871c11 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Fri, 1 May 2020 14:52:02 -0400 Subject: [PATCH 20/49] Update intersphinx_mapping from PlasmaPy In order to get links to different packages in the docs. --- docs/conf.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/docs/conf.py b/docs/conf.py index 6727c8a..e735f81 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -55,7 +55,12 @@ # -- Options for intersphinx extension --------------------------------------- # Example configuration for intersphinx: refer to the Python standard library. -intersphinx_mapping = {"https://docs.python.org/": None} +intersphinx_mapping = { + 'python': ('https://docs.python.org/3', None), + 'numpy': ('https://docs.scipy.org/doc/numpy', None), + 'scipy': ('https://docs.scipy.org/doc/scipy/reference/', None), + 'pandas': ('http://pandas.pydata.org/pandas-docs/stable/', None), + 'astropy': ('http://docs.astropy.org/en/stable/', None)} # -- Options for HTML output ------------------------------------------------- From 74f464bb34a03e1b23535bea4913505b1ac4b726 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Fri, 1 May 2020 14:52:40 -0400 Subject: [PATCH 21/49] Begin adding parameters to docstring for EigenData --- plasmapy_nei/eigen/eigenclass.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/plasmapy_nei/eigen/eigenclass.py b/plasmapy_nei/eigen/eigenclass.py index 43b4919..653b1b8 100644 --- a/plasmapy_nei/eigen/eigenclass.py +++ b/plasmapy_nei/eigen/eigenclass.py @@ -26,6 +26,12 @@ def _get_equilibrium_charge_states(ioniz_rate, recomb_rate, natom): Parameters ---------- + ioniz_rate + + recomb_rate + + natom + The atomic number. """ nstates = natom + 1 conce = np.zeros(nstates) From b4eee780e71dfba8eec0fab086281de01e5551eb Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Fri, 1 May 2020 14:53:09 -0400 Subject: [PATCH 22/49] Fix capitalization error --- plasmapy_nei/nei/nei.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/plasmapy_nei/nei/nei.py b/plasmapy_nei/nei/nei.py index 83c2988..f06780c 100644 --- a/plasmapy_nei/nei/nei.py +++ b/plasmapy_nei/nei/nei.py @@ -35,6 +35,7 @@ class NEIError(Exception): + """For when there are errors in setting up or performing NEI simulations.""" pass @@ -49,10 +50,10 @@ class SimulationResults: states of different elements and plasma properties as the initial conditions. - n_init: astropy.units.quantity + n_init: astropy.units.Quantity The initial number density scaling factor. - T_e_init: astropy.units.quantity + T_e_init: astropy.units.Quantity The initial electron temperature. max_steps: int From ba3edbc69dc440f22fdeb96185a465b1a6beeb8a Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Fri, 1 May 2020 15:56:37 -0400 Subject: [PATCH 23/49] Clarify conditional statement Co-authored-by: Erik Everson --- plasmapy_nei/eigen/eigenclass.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plasmapy_nei/eigen/eigenclass.py b/plasmapy_nei/eigen/eigenclass.py index 653b1b8..dd591ad 100644 --- a/plasmapy_nei/eigen/eigenclass.py +++ b/plasmapy_nei/eigen/eigenclass.py @@ -304,7 +304,7 @@ def eigenvalues(self, T_e=None, T_e_index=None): Returns the eigenvalues for the ionization and recombination rates for the temperature specified in the class. """ - if T_e_index: + if T_e_index is not None: return self._eigenvalues[T_e_index, :] elif T_e: T_e_index = self._get_temperature_index(T_e) From 89ee09a674b2d7c0e152094281a4b8ae703822ab Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Fri, 1 May 2020 15:57:31 -0400 Subject: [PATCH 24/49] Add PlasmaPy to intersphinx mapping Co-authored-by: Erik Everson --- docs/conf.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/conf.py b/docs/conf.py index e735f81..beb1dfb 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -60,7 +60,9 @@ 'numpy': ('https://docs.scipy.org/doc/numpy', None), 'scipy': ('https://docs.scipy.org/doc/scipy/reference/', None), 'pandas': ('http://pandas.pydata.org/pandas-docs/stable/', None), - 'astropy': ('http://docs.astropy.org/en/stable/', None)} + 'astropy': ('http://docs.astropy.org/en/stable/', None), + 'plasmapy': ('http://docs.plasmapy.org/en/latest/', None), +} # -- Options for HTML output ------------------------------------------------- From 3133cfff6b7cb35daa282d14c38a9c4ec43f3ca6 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Fri, 1 May 2020 16:20:50 -0400 Subject: [PATCH 25/49] Use `size` attribute instead of calling `len` Co-authored-by: Erik Everson --- plasmapy_nei/eigen/eigenclass.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plasmapy_nei/eigen/eigenclass.py b/plasmapy_nei/eigen/eigenclass.py index dd591ad..a252918 100644 --- a/plasmapy_nei/eigen/eigenclass.py +++ b/plasmapy_nei/eigen/eigenclass.py @@ -115,7 +115,7 @@ def _load_data(self): with h5py.File(filename, "r") as file: self._temperature_grid = file["te_gird"][:] # TODO: fix typo in HDF5 file - self._ntemp = len(self.temperature_grid) + self._ntemp = self._temperature_grid.size c_ori = file["ioniz_rate"][:] r_ori = file["recomb_rate"][:] From c3e49fbae31b5d96de67a14afaa3e1f840ead7b0 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Fri, 1 May 2020 16:26:46 -0400 Subject: [PATCH 26/49] Specify minimum package versions Co-authored-by: Erik Everson --- setup.cfg | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/setup.cfg b/setup.cfg index e261197..91712e5 100644 --- a/setup.cfg +++ b/setup.cfg @@ -15,9 +15,10 @@ python_requires = >=3.7 setup_requires = setuptools_scm install_requires = - astropy - plasmapy + astropy>=3.2 h5py + numpy>=1.16 + plasmapy>=0.3.1 [options.entry_points] @@ -120,4 +121,3 @@ max-line-length = 88 line_length=88 multi_line_output=3 include_trailing_comma: True - From 1b7460e6d13650c594c3bdb1cc8c7dcea9877b9b Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Fri, 1 May 2020 16:30:49 -0400 Subject: [PATCH 27/49] Be more specific in a conditional Co-authored-by: Erik Everson --- plasmapy_nei/eigen/eigenclass.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plasmapy_nei/eigen/eigenclass.py b/plasmapy_nei/eigen/eigenclass.py index a252918..4242caf 100644 --- a/plasmapy_nei/eigen/eigenclass.py +++ b/plasmapy_nei/eigen/eigenclass.py @@ -306,7 +306,7 @@ def eigenvalues(self, T_e=None, T_e_index=None): """ if T_e_index is not None: return self._eigenvalues[T_e_index, :] - elif T_e: + elif T_e is not None: T_e_index = self._get_temperature_index(T_e) return self._eigenvalues[T_e_index, :] elif self.temperature: From 8bb9cabc7e3843b007fe623a88f3eb1b5df92773 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Fri, 1 May 2020 16:31:56 -0400 Subject: [PATCH 28/49] Remove install_requires in setup.py This should be handled by within setup.cfg. Co-authored-by: Erik Everson --- setup.py | 1 - 1 file changed, 1 deletion(-) diff --git a/setup.py b/setup.py index 6fdd635..655bdf8 100755 --- a/setup.py +++ b/setup.py @@ -20,6 +20,5 @@ "write_to": os.path.join("plasmapy_nei", "version.py"), "write_to_template": VERSION_TEMPLATE, }, - install_requires=["plasmapy>=0.3.1", "numpy>=1.16", "astropy>=3.2", "h5py"], include_package_data=True, ) From 87415c18040e1dfc9bbfb198b036ca395b831500 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Fri, 1 May 2020 18:51:33 -0400 Subject: [PATCH 29/49] Limit sphinx to 2.4.4 until automodapi fix is released --- setup.cfg | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.cfg b/setup.cfg index e261197..e4d2d5d 100644 --- a/setup.cfg +++ b/setup.cfg @@ -32,7 +32,7 @@ test = pytest-doctestplus pytest-cov docs = - sphinx + sphinx <= 2.4.4 sphinx-automodapi [options.package_data] From a56591c24edb13086faac52bec6ee684e9c2643e Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Fri, 1 May 2020 18:52:25 -0400 Subject: [PATCH 30/49] Minor updates to __init__.py --- plasmapy_nei/__init__.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/plasmapy_nei/__init__.py b/plasmapy_nei/__init__.py index 4670952..2872ee5 100644 --- a/plasmapy_nei/__init__.py +++ b/plasmapy_nei/__init__.py @@ -1,11 +1,11 @@ """A Python package for non-equilibrium ionization modelingg of plasma.""" +__all__ = ["eigen", "nei"] + import warnings try: from .version import __version__ - - del version except Exception as exc: warnings.warn("Unable to import __version__") finally: @@ -13,5 +13,3 @@ from . import eigen from . import nei - -__all__ = ["eigen", "nei"] From 34c2db727a2a1b58beef6fd9d52128fa39c12c77 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Mon, 4 May 2020 12:50:48 -0400 Subject: [PATCH 31/49] Reformat code with black --- docs/conf.py | 12 ++++++------ plasmapy_nei/eigen/eigenclass.py | 7 +++++-- plasmapy_nei/nei/nei.py | 1 + 3 files changed, 12 insertions(+), 8 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index beb1dfb..f21efec 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -56,12 +56,12 @@ # Example configuration for intersphinx: refer to the Python standard library. intersphinx_mapping = { - 'python': ('https://docs.python.org/3', None), - 'numpy': ('https://docs.scipy.org/doc/numpy', None), - 'scipy': ('https://docs.scipy.org/doc/scipy/reference/', None), - 'pandas': ('http://pandas.pydata.org/pandas-docs/stable/', None), - 'astropy': ('http://docs.astropy.org/en/stable/', None), - 'plasmapy': ('http://docs.plasmapy.org/en/latest/', None), + "python": ("https://docs.python.org/3", None), + "numpy": ("https://docs.scipy.org/doc/numpy", None), + "scipy": ("https://docs.scipy.org/doc/scipy/reference/", None), + "pandas": ("http://pandas.pydata.org/pandas-docs/stable/", None), + "astropy": ("http://docs.astropy.org/en/stable/", None), + "plasmapy": ("http://docs.plasmapy.org/en/latest/", None), } # -- Options for HTML output ------------------------------------------------- diff --git a/plasmapy_nei/eigen/eigenclass.py b/plasmapy_nei/eigen/eigenclass.py index 4242caf..2bddf97 100644 --- a/plasmapy_nei/eigen/eigenclass.py +++ b/plasmapy_nei/eigen/eigenclass.py @@ -33,6 +33,9 @@ def _get_equilibrium_charge_states(ioniz_rate, recomb_rate, natom): natom The atomic number. """ + + # TODO: refactor this function with + nstates = natom + 1 conce = np.zeros(nstates) f = np.zeros(nstates + 1) @@ -47,7 +50,7 @@ def _get_equilibrium_charge_states(ioniz_rate, recomb_rate, natom): f[1] = 1.0 f[2] = c[1] * f[1] / r[2] - # simplest for hydrogen + # The solution for hydrogen may be found analytically. if natom == 1: f[1] = 1.0 / (1.0 + c[1] / r[2]) f[2] = c[1] * f[1] / r[2] @@ -70,7 +73,7 @@ def _get_equilibrium_charge_states(ioniz_rate, recomb_rate, natom): f[natom + 1] = c[natom] * f[natom] / r[natom + 1] - conce[0:nstates] = f[1 : nstates + 1] + conce[0:nstates] = f[1: nstates + 1] return conce diff --git a/plasmapy_nei/nei/nei.py b/plasmapy_nei/nei/nei.py index f06780c..6f69635 100644 --- a/plasmapy_nei/nei/nei.py +++ b/plasmapy_nei/nei/nei.py @@ -36,6 +36,7 @@ class NEIError(Exception): """For when there are errors in setting up or performing NEI simulations.""" + pass From 94941441bf2d40f86f35379a832c2bb3165f799a Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Mon, 4 May 2020 12:51:57 -0400 Subject: [PATCH 32/49] Formatting changes --- plasmapy_nei/nei/nei.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/plasmapy_nei/nei/nei.py b/plasmapy_nei/nei/nei.py index 6f69635..065f238 100644 --- a/plasmapy_nei/nei/nei.py +++ b/plasmapy_nei/nei/nei.py @@ -1112,7 +1112,7 @@ def _set_adaptive_timestep(self): # Find the boundaries to the temperature bin. index = self._get_temperature_index(T.to(u.K).value) - T_nearby = np.array(self._temperature_grid[index - 1 : index + 2]) * u.K + T_nearby = np.array(self._temperature_grid[index - 1: index + 2]) * u.K T_boundary = (T_nearby[0:-1] + T_nearby[1:]) / 2 # In order to use Brent's method, we must bound the root's @@ -1159,7 +1159,7 @@ def _set_adaptive_timestep(self): # and after the temperature leaves the temperature bin as bounds # for the root finding method. - dt_bounds = (dt_spread[first_false_index - 1 : first_false_index + 1]).value + dt_bounds = (dt_spread[first_false_index - 1: first_false_index + 1]).value # Define a function for the difference between the temperature # and the temperature boundary as a function of the value of the From add4eea7cac2d57df72b886b786006ab792e9e87 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dominik=20Sta=C5=84czak?= Date: Fri, 8 May 2020 16:16:35 +0200 Subject: [PATCH 33/49] Make code safe from atomic -> particles transition --- plasmapy_nei/nei/nei.py | 7 +++++-- plasmapy_nei/nei/tests/test_nei.py | 5 ++++- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/plasmapy_nei/nei/nei.py b/plasmapy_nei/nei/nei.py index f06780c..c300989 100644 --- a/plasmapy_nei/nei/nei.py +++ b/plasmapy_nei/nei/nei.py @@ -9,7 +9,10 @@ import plasmapy as pl from scipy import interpolate, optimize from plasmapy_nei.eigen import EigenData -from plasmapy.atomic import IonizationStates +try: + from plasmapy.atomic import IonizationStates, atomic_number +except ImportError: + from plasmapy.particles import IonizationStates, atomic_number import warnings @@ -79,7 +82,7 @@ def __init__( self._max_steps = max_steps self._nstates = { - elem: pl.atomic.atomic_number(elem) + 1 for elem in self.elements + elem: atomic_number(elem) + 1 for elem in self.elements } self._ionic_fractions = { diff --git a/plasmapy_nei/nei/tests/test_nei.py b/plasmapy_nei/nei/tests/test_nei.py index 1f53944..ff1be39 100644 --- a/plasmapy_nei/nei/tests/test_nei.py +++ b/plasmapy_nei/nei/tests/test_nei.py @@ -1,7 +1,10 @@ """Tests for non-equilibrium ionization modeling classes.""" import astropy.units as u -from plasmapy.atomic import IonizationStates, particle_symbol +try: + from plasmapy.atomic import IonizationStates, particle_symbol +except ImportError: + from plasmapy.particles import IonizationStates, particle_symbol from plasmapy_nei.nei import NEI from plasmapy_nei.eigen import EigenData import numpy as np From 3dce3f6d8f9fdadaba5a262f77de45318ef96631 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dominik=20Sta=C5=84czak?= Date: Fri, 8 May 2020 16:17:07 +0200 Subject: [PATCH 34/49] Initialize _results with None by default It's bad practice to create named entities for the class in functions elsewhere - it means you can't just inspect the __init__ looking for variables. --- plasmapy_nei/nei/nei.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/plasmapy_nei/nei/nei.py b/plasmapy_nei/nei/nei.py index c300989..35b7a60 100644 --- a/plasmapy_nei/nei/nei.py +++ b/plasmapy_nei/nei/nei.py @@ -509,6 +509,8 @@ def __init__( self._get_temperature_index = self._EigenDataDict[ self.elements[0] ]._get_temperature_index + + self._results = None except Exception: raise NEIError( @@ -1000,9 +1002,9 @@ def results(self) -> SimulationResults: corresponds to the simulation results. """ - try: + if self._results is not None: return self._results - except Exception: + else: raise AttributeError("The simulation has not yet been performed.") @property From 9f484a1ed3c56c4bfab3849258d07cee9cc70e11 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dominik=20Sta=C5=84czak?= Date: Fri, 8 May 2020 16:19:08 +0200 Subject: [PATCH 35/49] Raise __init__ NEIError from the original exception because losing the whole traceback seems unhelpful for debugging why the traceback is popping up! --- plasmapy_nei/nei/nei.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/plasmapy_nei/nei/nei.py b/plasmapy_nei/nei/nei.py index 35b7a60..17fe827 100644 --- a/plasmapy_nei/nei/nei.py +++ b/plasmapy_nei/nei/nei.py @@ -512,7 +512,7 @@ def __init__( self._results = None - except Exception: + except Exception as e: raise NEIError( f"Unable to create NEI instance for:\n" f" inputs = {inputs}\n" @@ -523,7 +523,7 @@ def __init__( f" time_start = {time_start}\n" f" time_max = {time_max}\n" f" max_steps = {max_steps}\n" - ) + ) from e def equil_ionic_fractions( self, T_e: u.Quantity = None, time: u.Quantity = None, From b87cb890f46cc2750c42b144d4367490bc54d648 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dominik=20Sta=C5=84czak?= Date: Fri, 8 May 2020 16:20:08 +0200 Subject: [PATCH 36/49] Rewrite test_nei using fixtures This is meant to clean up the code using fixtures for best practices. Nick, I really hope you read this one ;) The idea is: * we get rid of the class - there are no other tests than for the NEI object in the entire file, so they're already grouped up nicely * we replace the three objects: list of test case names, list of test input dicts and (class-stored) list of initialized NEI instances with a fixture containing the three elements * we can thus remove most of the @pytest.mark.parametrize entries and call the fixture by its name - we do have to unwrap the tuple into the three objects in each tests, but I haven't yet figured out how else to do that * we remove the initialization test as the fixture itself serves that purpose * we concatenate test_final_results and test_simulate because we don't want to depend on the order in which the tests are ran --- plasmapy_nei/nei/tests/test_nei.py | 311 +++++++++++++---------------- 1 file changed, 143 insertions(+), 168 deletions(-) diff --git a/plasmapy_nei/nei/tests/test_nei.py b/plasmapy_nei/nei/tests/test_nei.py index ff1be39..eca27b7 100644 --- a/plasmapy_nei/nei/tests/test_nei.py +++ b/plasmapy_nei/nei/tests/test_nei.py @@ -18,8 +18,9 @@ n_array = np.array([1e9, 5e8]) * u.cm ** -3 -tests = { - "basic": { +@pytest.fixture(scope="module", params = + [ + ("basic", { "inputs": inputs_dict, "abundances": abundances, "T_e": T_e_array, @@ -31,8 +32,8 @@ "dt": 800 * u.s, "adapt_dt": False, "verbose": True, - }, - "T_e constant": { + }), + ("T_e constant", { "inputs": inputs_dict, "abundances": abundances, "T_e": 1 * u.MK, @@ -44,8 +45,8 @@ "max_steps": 2, "adapt_dt": False, "verbose": True, - }, - "n_e constant": { + }), + ("n_e constant", { "inputs": inputs_dict, "abundances": abundances, "T_e": T_e_array, @@ -57,8 +58,8 @@ "adapt_dt": False, "dt": 100 * u.s, "verbose": True, - }, - "T_e function": { + }), + ("T_e function", { "inputs": inputs_dict, "abundances": abundances, "T_e": lambda time: 1e4 * (1 + time / u.s) * u.K, @@ -68,8 +69,8 @@ "dt": 100 * u.s, "adapt_dt": False, "verbose": True, - }, - "n function": { + }), + ("n function", { "inputs": inputs_dict, "abundances": abundances, "T_e": 6e4 * u.K, @@ -80,8 +81,8 @@ "dt": 200 * u.s, "verbose": True, "max_steps": 4, - }, - "equil test cool": { + }), + ("equil test cool", { "inputs": ["H", "He", "N"], "abundances": {"H": 1, "He": 0.1, "C": 1e-4, "N": 1e-4, "O": 1e-4, "Fe": 1e-4}, "T_e": 10001.0 * u.K, @@ -92,8 +93,8 @@ "dt": 5e5 * u.s, "max_steps": 4, "verbose": True, - }, - "equil test hot": { + }), + ("equil test hot", { "inputs": ["H", "He", "C"], "abundances": { "H": 1, @@ -111,8 +112,8 @@ "max_steps": 3, "adapt_dt": False, "verbose": True, - }, - "equil test start far out of equil": { + }), + ("equil test start far out of equil", { "inputs": { "H": [0.99, 0.01], "He": [0.5, 0.0, 0.5], @@ -127,8 +128,8 @@ "adapt_dt": False, "verbose": True, "max_steps": 2, - }, - "adapt dt": { + }), + ("adapt dt", { "inputs": ["H", "He"], "abundances": {"H": 1, "He": 0.1}, "T_e": lambda t: u.K * (1e6 + 1.3e4 * np.sin(t.value)), @@ -137,173 +138,147 @@ "time_start": 0 * u.s, "time_max": 2 * np.pi * u.s, "adapt_dt": True, - }, -} + }), +]) +def name_inputs_instance(request): + test_name, dictionary = request.param + return test_name, dictionary, NEI(**dictionary) -test_names = list(tests.keys()) +def test_time_start(name_inputs_instance): + test_name, inputs, instance = name_inputs_instance + if "time_start" in inputs.keys(): + assert inputs["time_start"] == instance.time_start + elif "time_input" in inputs.keys(): + assert inputs["time_input"].min() == instance.time_start + else: + assert instance.time_start == 0 * u.s -class TestNEI: - @classmethod - def setup_class(cls): - cls.instances = {} +def test_time_max(name_inputs_instance): + test_name, inputs, instance = name_inputs_instance + if "time_max" in inputs.keys(): + assert inputs["time_max"] == instance.time_max - @pytest.mark.parametrize("test_name", test_names) - def test_instantiate(self, test_name): - try: - instance = NEI(**tests[test_name]) - self.instances[test_name] = instance - except Exception as exc: - raise Exception(f"Problem with test {test_name}") from exc +def test_initial_type(name_inputs_instance): + test_name, inputs, instance = name_inputs_instance + assert isinstance(instance.initial, IonizationStates) - @pytest.mark.parametrize("test_name", test_names) - def test_time_start(self, test_name): - instance = self.instances[test_name] - if "time_start" in tests[test_name].keys(): - assert tests[test_name]["time_start"] == instance.time_start - elif "time_input" in tests[test_name].keys(): - assert tests[test_name]["time_input"].min() == instance.time_start - else: - assert instance.time_start == 0 * u.s - - @pytest.mark.parametrize("test_name", test_names) - def test_time_max(self, test_name): - instance = self.instances[test_name] - if "time_max" in tests[test_name].keys(): - assert tests[test_name]["time_max"] == instance.time_max +def test_n_input(name_inputs_instance): + test_name, inputs, instance = name_inputs_instance + actual = instance.n_input + expected = inputs["n"] + if isinstance(expected, u.Quantity) and not expected.isscalar: + assert all(expected == actual) + else: + assert expected == actual - @pytest.mark.parametrize("test_name", test_names) - def test_initial_type(self, test_name): - instance = self.instances[test_name] - assert isinstance(instance.initial, IonizationStates) +def test_T_e_input(name_inputs_instance): + test_name, inputs, instance = name_inputs_instance + actual = instance.T_e_input + expected = inputs["T_e"] + if isinstance(expected, u.Quantity) and not expected.isscalar: + assert all(expected == actual) + else: + assert expected == actual - @pytest.mark.parametrize("test_name", test_names) - def test_n_input(self, test_name): - actual = self.instances[test_name].n_input - expected = tests[test_name]["n"] - if isinstance(expected, u.Quantity) and not expected.isscalar: - assert all(expected == actual) +def test_electron_temperature(name_inputs_instance): + test_name, inputs, instance = name_inputs_instance + instance = instance + T_e_input = instance.T_e_input + if isinstance(T_e_input, u.Quantity): + if T_e_input.isscalar: + assert instance.electron_temperature(instance.time_start) == T_e_input + assert instance.electron_temperature(instance.time_max) == T_e_input else: - assert expected == actual - - @pytest.mark.parametrize("test_name", test_names) - def test_T_e_input(self, test_name): - actual = self.instances[test_name].T_e_input - expected = tests[test_name]["T_e"] - if isinstance(expected, u.Quantity) and not expected.isscalar: - assert all(expected == actual) - else: - assert expected == actual - - @pytest.mark.parametrize("test_name", test_names) - def test_electron_temperature(self, test_name): - instance = self.instances[test_name] - T_e_input = instance.T_e_input - if isinstance(T_e_input, u.Quantity): - if T_e_input.isscalar: - assert instance.electron_temperature(instance.time_start) == T_e_input - assert instance.electron_temperature(instance.time_max) == T_e_input - else: - for time, T_e in zip(instance.time_input, T_e_input): - try: - T_e_func = instance.electron_temperature(time) - except Exception: - raise ValueError("Unable to find T_e from electron_temperature") + for time, T_e in zip(instance.time_input, T_e_input): + try: + T_e_func = instance.electron_temperature(time) + except Exception: + raise ValueError("Unable to find T_e from electron_temperature") - assert np.isclose(T_e.value, T_e_func.value) - if callable(T_e_input): - assert instance.T_e_input( - instance.time_start - ) == instance.electron_temperature(instance.time_start) + assert np.isclose(T_e.value, T_e_func.value) + if callable(T_e_input): + assert instance.T_e_input( + instance.time_start + ) == instance.electron_temperature(instance.time_start) - @pytest.mark.parametrize( - "test_name", - [ - test_name - for test_name in test_names - if isinstance(tests[test_name]["inputs"], dict) - ], - ) - def test_initial_ionfracs(self, test_name): - original_inputs = tests[test_name]["inputs"] - original_elements = original_inputs.keys() +def test_initial_ionfracs(name_inputs_instance): + test_name, inputs, instance = name_inputs_instance + if not isinstance(inputs["inputs"], dict): + pytest.skip("Test irrelevant") - for element in original_elements: - assert np.allclose( - original_inputs[element], - self.instances[test_name].initial.ionic_fractions[ - particle_symbol(element) - ], - ) + + original_inputs = inputs["inputs"] + original_elements = original_inputs.keys() - @pytest.mark.parametrize("test_name", test_names) - def test_simulate(self, test_name): - try: - self.instances[test_name].simulate() - except Exception as exc: - raise ValueError(f"Unable to simulate for test: {test_name}") from exc + for element in original_elements: + assert np.allclose( + original_inputs[element], + instance.initial.ionic_fractions[ + particle_symbol(element) + ], + ) - @pytest.mark.parametrize("test_name", test_names) - def test_simulation_end(self, test_name): - instance = self.instances[test_name] - time = instance.results.time - end_time = time[-1] - max_steps = instance.max_steps +def test_simulate(name_inputs_instance): + test_name, inputs, instance = name_inputs_instance + instance.simulate() + final = instance.final + results = instance.results + # Make sure the elements are equal to each other + assert final.ionic_fractions.keys() == results.ionic_fractions.keys() + assert final.abundances == results.abundances + for elem in final.ionic_fractions.keys(): + assert np.allclose( + results.ionic_fractions[elem][-1, :], final.ionic_fractions[elem] + ) - if np.isnan(end_time.value): - raise Exception("End time is NaN.") +def test_simulation_end(name_inputs_instance): + test_name, inputs, instance = name_inputs_instance + instance = instance + time = instance.results.time + end_time = time[-1] + max_steps = instance.max_steps - got_to_max_steps = len(time) == instance.max_steps + 1 - got_to_time_max = np.isclose(time[-1].value, instance.time_max.value) + if np.isnan(end_time.value): + raise Exception("End time is NaN.") - if time.isscalar or len(time) == 1: - raise Exception(f"The only element in results.time is {time}") + got_to_max_steps = len(time) == instance.max_steps + 1 + got_to_time_max = np.isclose(time[-1].value, instance.time_max.value) - if not got_to_max_steps and not got_to_time_max: - print(f"time = {time}") - print(f"max_steps = {max_steps}") - print(f"time_max = {instance.time_max}") - raise Exception("Problem with end time.") + if time.isscalar or len(time) == 1: + raise Exception(f"The only element in results.time is {time}") - @pytest.mark.parametrize( - "test_name", [test_name for test_name in test_names if "equil" in test_name] - ) - def test_equilibration(self, test_name): - """ - Test that equilibration works. - """ - instance = self.instances[test_name] - T_e = instance.T_e_input - assert ( - isinstance(T_e, u.Quantity) and T_e.isscalar - ), "This test can only be used for cases where T_e is constant." - equil_dict = instance.equil_ionic_fractions(T_e) - for element in instance.elements: - assert np.allclose( - equil_dict[element], instance.results.ionic_fractions[element][-1, :] - ) + if not got_to_max_steps and not got_to_time_max: + print(f"time = {time}") + print(f"max_steps = {max_steps}") + print(f"time_max = {instance.time_max}") + raise Exception("Problem with end time.") - @pytest.mark.parametrize("test_name", test_names) - def test_initial_results(self, test_name): - initial = self.instances[test_name].initial - results = self.instances[test_name].results - # Make sure that the elements are equal to each other - assert initial.ionic_fractions.keys() == results.ionic_fractions.keys() - assert initial.abundances == results.abundances - for elem in initial.ionic_fractions.keys(): # TODO: enable initial.elements - assert np.allclose( - results.ionic_fractions[elem][0, :], initial.ionic_fractions[elem] - ) +def test_equilibration(name_inputs_instance): + test_name, inputs, instance = name_inputs_instance + if "equil" in test_name: + pytest.skip("Test irrelevant") + """ + Test that equilibration works. + """ + instance = instance + T_e = instance.T_e_input + if not (isinstance(T_e, u.Quantity) and T_e.isscalar): + pytest.skip("This test can only be used for cases where T_e is constant.") + equil_dict = instance.equil_ionic_fractions(T_e) + for element in instance.elements: + assert np.allclose( + equil_dict[element], instance.results.ionic_fractions[element][-1, :] + ) - @pytest.mark.parametrize("test_name", test_names) - def test_final_results(self, test_name): - # initial = self.instances[test_name].initial - final = self.instances[test_name].final - results = self.instances[test_name].results - # Make sure the elements are equal to each other - assert final.ionic_fractions.keys() == results.ionic_fractions.keys() - assert final.abundances == results.abundances - for elem in final.ionic_fractions.keys(): - assert np.allclose( - results.ionic_fractions[elem][-1, :], final.ionic_fractions[elem] - ) +def test_initial_results(name_inputs_instance): + test_name, inputs, instance = name_inputs_instance + initial = instance.initial + results = instance.results + # Make sure that the elements are equal to each other + assert initial.ionic_fractions.keys() == results.ionic_fractions.keys() + assert initial.abundances == results.abundances + for elem in initial.ionic_fractions.keys(): # TODO: enable initial.elements + assert np.allclose( + results.ionic_fractions[elem][0, :], initial.ionic_fractions[elem] + ) From b5acd1cdfb0ebeb57baa782bb869cadad8bfebd6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dominik=20Sta=C5=84czak?= Date: Fri, 8 May 2020 16:30:11 +0200 Subject: [PATCH 37/49] Refactor: use numpy to simplify comparison --- plasmapy_nei/nei/tests/test_nei.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/plasmapy_nei/nei/tests/test_nei.py b/plasmapy_nei/nei/tests/test_nei.py index eca27b7..eb76bb1 100644 --- a/plasmapy_nei/nei/tests/test_nei.py +++ b/plasmapy_nei/nei/tests/test_nei.py @@ -167,19 +167,13 @@ def test_n_input(name_inputs_instance): test_name, inputs, instance = name_inputs_instance actual = instance.n_input expected = inputs["n"] - if isinstance(expected, u.Quantity) and not expected.isscalar: - assert all(expected == actual) - else: - assert expected == actual + assert np.all(expected == actual) def test_T_e_input(name_inputs_instance): test_name, inputs, instance = name_inputs_instance actual = instance.T_e_input expected = inputs["T_e"] - if isinstance(expected, u.Quantity) and not expected.isscalar: - assert all(expected == actual) - else: - assert expected == actual + assert np.all(expected == actual) def test_electron_temperature(name_inputs_instance): test_name, inputs, instance = name_inputs_instance From bb8cc0212160af1a3b9e9e19a609abdc6ead4774 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dominik=20Sta=C5=84czak?= Date: Fri, 8 May 2020 16:31:56 +0200 Subject: [PATCH 38/49] Refactor: remove uninformative exception --- plasmapy_nei/nei/tests/test_nei.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/plasmapy_nei/nei/tests/test_nei.py b/plasmapy_nei/nei/tests/test_nei.py index eb76bb1..3fda382 100644 --- a/plasmapy_nei/nei/tests/test_nei.py +++ b/plasmapy_nei/nei/tests/test_nei.py @@ -185,10 +185,7 @@ def test_electron_temperature(name_inputs_instance): assert instance.electron_temperature(instance.time_max) == T_e_input else: for time, T_e in zip(instance.time_input, T_e_input): - try: - T_e_func = instance.electron_temperature(time) - except Exception: - raise ValueError("Unable to find T_e from electron_temperature") + T_e_func = instance.electron_temperature(time) assert np.isclose(T_e.value, T_e_func.value) if callable(T_e_input): From 7be59d8c48755d56bbdeec2b6c9b7dcf4de783b2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dominik=20Sta=C5=84czak?= Date: Fri, 8 May 2020 16:32:49 +0200 Subject: [PATCH 39/49] Refactor: a silly autoreplace bug --- plasmapy_nei/nei/tests/test_nei.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/plasmapy_nei/nei/tests/test_nei.py b/plasmapy_nei/nei/tests/test_nei.py index 3fda382..90ef7fe 100644 --- a/plasmapy_nei/nei/tests/test_nei.py +++ b/plasmapy_nei/nei/tests/test_nei.py @@ -177,7 +177,6 @@ def test_T_e_input(name_inputs_instance): def test_electron_temperature(name_inputs_instance): test_name, inputs, instance = name_inputs_instance - instance = instance T_e_input = instance.T_e_input if isinstance(T_e_input, u.Quantity): if T_e_input.isscalar: @@ -225,7 +224,6 @@ def test_simulate(name_inputs_instance): def test_simulation_end(name_inputs_instance): test_name, inputs, instance = name_inputs_instance - instance = instance time = instance.results.time end_time = time[-1] max_steps = instance.max_steps @@ -252,7 +250,6 @@ def test_equilibration(name_inputs_instance): """ Test that equilibration works. """ - instance = instance T_e = instance.T_e_input if not (isinstance(T_e, u.Quantity) and T_e.isscalar): pytest.skip("This test can only be used for cases where T_e is constant.") From ef34433d7cf07e97252c4901c908abcf48c16271 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dominik=20Sta=C5=84czak?= Date: Fri, 8 May 2020 16:40:32 +0200 Subject: [PATCH 40/49] Remove unnecessary import --- plasmapy_nei/nei/nei.py | 1 - 1 file changed, 1 deletion(-) diff --git a/plasmapy_nei/nei/nei.py b/plasmapy_nei/nei/nei.py index 17fe827..0153845 100644 --- a/plasmapy_nei/nei/nei.py +++ b/plasmapy_nei/nei/nei.py @@ -6,7 +6,6 @@ import numpy as np from typing import Union, Optional, List, Dict, Callable import astropy.units as u -import plasmapy as pl from scipy import interpolate, optimize from plasmapy_nei.eigen import EigenData try: From 00bf381d8df91d5c2e5e8eb907c1fb7cf40cb4f4 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Fri, 8 May 2020 23:05:34 -0400 Subject: [PATCH 41/49] Use slightly faster way of constructing a Quantity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Dominik Stańczak --- plasmapy_nei/nei/nei.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plasmapy_nei/nei/nei.py b/plasmapy_nei/nei/nei.py index 0153845..d43e457 100644 --- a/plasmapy_nei/nei/nei.py +++ b/plasmapy_nei/nei/nei.py @@ -395,7 +395,7 @@ class NEI: >>> inputs = {'H': [0.9, 0.1], 'He': [0.9, 0.099, 0.001]} >>> abund = {'H': 1, 'He': 0.085} - >>> n = np.array([1e9, 1e8]) * u.cm ** -3 + >>> n = u.Quantity([1e9, 1e8], u.cm**-3) >>> T_e = np.array([10000, 40000]) * u.K >>> time = np.array([0, 300]) * u.s >>> dt = 0.25 * u.s From e9ec8a9a0600aab96232763a9ae827c74c16ecd8 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Fri, 8 May 2020 23:08:47 -0400 Subject: [PATCH 42/49] Shorten calculations for equilibrium Co-authored-by: Erik Everson --- plasmapy_nei/eigen/eigenclass.py | 30 ++++++++---------------------- 1 file changed, 8 insertions(+), 22 deletions(-) diff --git a/plasmapy_nei/eigen/eigenclass.py b/plasmapy_nei/eigen/eigenclass.py index 4242caf..860c1d4 100644 --- a/plasmapy_nei/eigen/eigenclass.py +++ b/plasmapy_nei/eigen/eigenclass.py @@ -44,31 +44,17 @@ def _get_equilibrium_charge_states(ioniz_rate, recomb_rate, natom): c[i + 1] = ioniz_rate[i] r[i + 1] = recomb_rate[i] + # f[0] = 0.0 from initialization f[1] = 1.0 - f[2] = c[1] * f[1] / r[2] - - # simplest for hydrogen - if natom == 1: - f[1] = 1.0 / (1.0 + c[1] / r[2]) - f[2] = c[1] * f[1] / r[2] - conce[0:2] = f[1:3] - return conce - - # for other elements - - for k in range(2, natom): - f[k + 1] = (-c[k - 1] * f[k - 1] + (c[k] + r[k]) * f[k]) / r[k + 1] - - f[natom + 1] = c[natom] * f[natom] / r[natom + 1] - - f[1] = 1.0 / np.sum(f) - - f[2] = c[1] * f[1] / r[2] - - for k in range(2, natom): - f[k + 1] = (-c[k - 1] * f[k - 1] + (c[k] + r[k]) * f[k]) / r[k + 1] + for k in range(2, natom+1): + # note: this loop does NOT run for the hydrogen case, thus it does note defined f[2] + f[k] = (-c[k - 2] / r[k]) * f[k - 2] \ + + ((c[k-1] + r[k-1]) / r[k]) * f[k-1]) f[natom + 1] = c[natom] * f[natom] / r[natom + 1] + + # normalize the distribution + f = f / np.sum(f) conce[0:nstates] = f[1 : nstates + 1] return conce From c790de33c0ff37a6b3f04df59ee3853388151352 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Mon, 1 Jun 2020 14:51:59 -0400 Subject: [PATCH 43/49] Change around imports for PlasmaPy 0.3.1 and after --- plasmapy_nei/eigen/tests/test_eigenclass.py | 6 +++++- plasmapy_nei/nei/nei.py | 2 ++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/plasmapy_nei/eigen/tests/test_eigenclass.py b/plasmapy_nei/eigen/tests/test_eigenclass.py index 18db8dd..18e7a8d 100644 --- a/plasmapy_nei/eigen/tests/test_eigenclass.py +++ b/plasmapy_nei/eigen/tests/test_eigenclass.py @@ -1,6 +1,10 @@ """Tests of the class to store package data.""" -from plasmapy import atomic +try: + from plasmapy import atomic +except ImportError: + from plasmapy import particles as atomic + from plasmapy_nei.eigen import EigenData import pytest import numpy as np diff --git a/plasmapy_nei/nei/nei.py b/plasmapy_nei/nei/nei.py index 8fc391e..c8ba909 100644 --- a/plasmapy_nei/nei/nei.py +++ b/plasmapy_nei/nei/nei.py @@ -8,10 +8,12 @@ import astropy.units as u from scipy import interpolate, optimize from plasmapy_nei.eigen import EigenData + try: from plasmapy.atomic import IonizationStates, atomic_number except ImportError: from plasmapy.particles import IonizationStates, atomic_number + import warnings From 31818e0522156c8743598512185602936fd34986 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Mon, 1 Jun 2020 20:50:32 -0400 Subject: [PATCH 44/49] Add exception handling directly to EigenData.__init__ --- plasmapy_nei/eigen/eigenclass.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/plasmapy_nei/eigen/eigenclass.py b/plasmapy_nei/eigen/eigenclass.py index f7a9e5f..8844315 100644 --- a/plasmapy_nei/eigen/eigenclass.py +++ b/plasmapy_nei/eigen/eigenclass.py @@ -228,8 +228,11 @@ def _load_data(self): @particle_input def __init__(self, element: Particle): - self._validate_element(element) - self._load_data() + try: + self._validate_element(element) + self._load_data() + except Exception as exc: + raise RuntimeError(f"Unable to create EigenData object for {element}") from exc def _get_temperature_index(self, T_e): # TODO: extract this to a function """Return the temperature index closest to a particular temperature.""" From 9c7fef9571812e9ca55855a4d94c8decf81293fa Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Wed, 3 Jun 2020 14:41:16 -0400 Subject: [PATCH 45/49] Add TODOs to eigenclass.py plus minor changes --- plasmapy_nei/eigen/eigenclass.py | 52 +++++++++++++++++++++++++++++--- 1 file changed, 47 insertions(+), 5 deletions(-) diff --git a/plasmapy_nei/eigen/eigenclass.py b/plasmapy_nei/eigen/eigenclass.py index 8844315..866b99d 100644 --- a/plasmapy_nei/eigen/eigenclass.py +++ b/plasmapy_nei/eigen/eigenclass.py @@ -27,14 +27,22 @@ def _get_equilibrium_charge_states(ioniz_rate, recomb_rate, natom): Parameters ---------- ioniz_rate + An array containing the ionization rates. recomb_rate + An array containing the recombination rates. natom The atomic number. """ - # TODO: refactor this function with + # TODO: specify what each index in the array means + + # TODO: possibly refactor to include data in form of xarray? + + # TODO: add notes to docs on the mathematics behind this calculation + + # TODO: use more descriptive variable names throughout function nstates = natom + 1 conce = np.zeros(nstates) @@ -74,11 +82,11 @@ def _get_equilibrium_charge_states(ioniz_rate, recomb_rate, natom): f[k + 1] = (-c[k - 1] * f[k - 1] + (c[k] + r[k]) * f[k]) / r[k + 1] f[natom + 1] = c[natom] * f[natom] / r[natom + 1] - + # normalize the distribution f = f / np.sum(f) - conce[0:nstates] = f[1: nstates + 1] + conce[0:nstates] = f[1 : nstates + 1] return conce @@ -232,7 +240,9 @@ def __init__(self, element: Particle): self._validate_element(element) self._load_data() except Exception as exc: - raise RuntimeError(f"Unable to create EigenData object for {element}") from exc + raise RuntimeError( + f"Unable to create EigenData object for {element}" + ) from exc def _get_temperature_index(self, T_e): # TODO: extract this to a function """Return the temperature index closest to a particular temperature.""" @@ -304,10 +314,12 @@ def ntemp(self) -> int: @property def ionization_rate(self): + # TODO: add docstring & description return self._ionization_rate @property def recombination_rate(self): + # TODO: add docstring & description return self._recombination_rate def eigenvalues(self, T_e=None, T_e_index=None): @@ -325,11 +337,23 @@ def eigenvalues(self, T_e=None, T_e_index=None): else: raise AttributeError("The temperature has not been set.") - def eigenvectors(self, T_e=None, T_e_index=None): + def eigenvectors(self, T_e: u.K = None, T_e_index: u.K = None): """ Returns the eigenvectors for the ionization and recombination rates for the temperature specified in the class. + + Parameters + ---------- + T_e : ~astropy.units.Quantity + The electron temperature + + T_e_index : integer + The index of the electron temperature array corresponding to + the desired temperature. """ + + # TODO: add discussion of what the indices of the returned array represent + if T_e_index: return self._eigenvectors[T_e_index, :, :] elif T_e: @@ -344,6 +368,15 @@ def eigenvector_inverses(self, T_e=None, T_e_index=None): """ Returns the inverses of the eigenvectors for the ionization and recombination rates for the temperature specified in the class. + + Parameters + ---------- + T_e : ~astropy.units.Quantity + The electron temperature + + T_e_index : integer + The index of the electron temperature array corresponding to + the desired temperature. """ if T_e_index: return self._eigenvector_inverses[T_e_index, :, :] @@ -359,6 +392,15 @@ def equilibrium_state(self, T_e=None, T_e_index=None): """ Return the equilibrium charge state distribution for the temperature specified in the class. + + Parameters + ---------- + T_e : ~astropy.units.Quantity + The electron temperature + + T_e_index : integer + The index of the electron temperature array corresponding to + the desired temperature. """ if T_e_index: return self._equilibrium_states[T_e_index, :] From 0d7d4425f2ac21907a64cade7e3bbf26a6c088b6 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Wed, 3 Jun 2020 15:13:27 -0400 Subject: [PATCH 46/49] Handle exceptions resulting from missing data file --- plasmapy_nei/eigen/eigenclass.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/plasmapy_nei/eigen/eigenclass.py b/plasmapy_nei/eigen/eigenclass.py index 866b99d..cda3a89 100644 --- a/plasmapy_nei/eigen/eigenclass.py +++ b/plasmapy_nei/eigen/eigenclass.py @@ -10,6 +10,7 @@ from numpy import linalg as LA import pkg_resources import warnings +import astropy.units as u try: from plasmapy.particles import Particle, particle_input @@ -127,13 +128,20 @@ def _load_data(self): "data/ionrecomb_rate.h5", # from Chianti database, version 8.07 ) - # TODO: enable different HDF5 files - - with h5py.File(filename, "r") as file: + try: + file = h5py.File(filename, "r") + except OSError as oserror: + raise IOError( + f"Unable to import {filename} using h5py. This error could " + f"happen, for example, if the repository was cloned without " + f"having git-lfs installed." + ) from oserror + else: self._temperature_grid = file["te_gird"][:] # TODO: fix typo in HDF5 file self._ntemp = self._temperature_grid.size c_ori = file["ioniz_rate"][:] r_ori = file["recomb_rate"][:] + file.close() c_rate = np.zeros((self.ntemp, self.nstates)) r_rate = np.zeros((self.ntemp, self.nstates)) From 7ce09579286aef258922d5366c318028016a4a41 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Wed, 3 Jun 2020 15:13:43 -0400 Subject: [PATCH 47/49] Clean up nei.py and tests --- plasmapy_nei/nei/nei.py | 16 +- plasmapy_nei/nei/tests/test_nei.py | 300 +++++++++++++++++------------ 2 files changed, 179 insertions(+), 137 deletions(-) diff --git a/plasmapy_nei/nei/nei.py b/plasmapy_nei/nei/nei.py index c8ba909..6b7e152 100644 --- a/plasmapy_nei/nei/nei.py +++ b/plasmapy_nei/nei/nei.py @@ -83,9 +83,7 @@ def __init__( self._abundances = initial.abundances self._max_steps = max_steps - self._nstates = { - elem: atomic_number(elem) + 1 for elem in self.elements - } + self._nstates = {elem: atomic_number(elem) + 1 for elem in self.elements} self._ionic_fractions = { elem: np.full((max_steps + 1, self.nstates[elem]), np.nan, dtype=np.float64) @@ -511,12 +509,12 @@ def __init__( self._get_temperature_index = self._EigenDataDict[ self.elements[0] ]._get_temperature_index - + self._results = None - except Exception as e: + except Exception as e: raise NEIError( - f"Unable to create NEI instance for:\n" + f"Unable to create NEI object for:\n" f" inputs = {inputs}\n" f" abundances = {abundances}\n" f" T_e = {T_e}\n" @@ -754,7 +752,7 @@ def dt_min(self, value: u.s): try: value = value.to(u.s) except u.UnitConversionError as exc: - raise u.UnitConversionError("Invalid units for dt_min.") + raise u.UnitConversionError("Invalid units for dt_min.") from exc if ( hasattr(self, "_dt_input") and self.dt_input is not None @@ -1118,7 +1116,7 @@ def _set_adaptive_timestep(self): # Find the boundaries to the temperature bin. index = self._get_temperature_index(T.to(u.K).value) - T_nearby = np.array(self._temperature_grid[index - 1: index + 2]) * u.K + T_nearby = np.array(self._temperature_grid[index - 1 : index + 2]) * u.K T_boundary = (T_nearby[0:-1] + T_nearby[1:]) / 2 # In order to use Brent's method, we must bound the root's @@ -1165,7 +1163,7 @@ def _set_adaptive_timestep(self): # and after the temperature leaves the temperature bin as bounds # for the root finding method. - dt_bounds = (dt_spread[first_false_index - 1: first_false_index + 1]).value + dt_bounds = (dt_spread[first_false_index - 1 : first_false_index + 1]).value # Define a function for the difference between the temperature # and the temperature boundary as a function of the value of the diff --git a/plasmapy_nei/nei/tests/test_nei.py b/plasmapy_nei/nei/tests/test_nei.py index 90ef7fe..8a947bd 100644 --- a/plasmapy_nei/nei/tests/test_nei.py +++ b/plasmapy_nei/nei/tests/test_nei.py @@ -1,6 +1,7 @@ """Tests for non-equilibrium ionization modeling classes.""" import astropy.units as u + try: from plasmapy.atomic import IonizationStates, particle_symbol except ImportError: @@ -18,128 +19,164 @@ n_array = np.array([1e9, 5e8]) * u.cm ** -3 -@pytest.fixture(scope="module", params = - [ - ("basic", { - "inputs": inputs_dict, - "abundances": abundances, - "T_e": T_e_array, - "n": n_array, - "time_input": time_array, - "time_start": 0 * u.s, - "time_max": 800 * u.s, - "max_steps": 1, - "dt": 800 * u.s, - "adapt_dt": False, - "verbose": True, - }), - ("T_e constant", { - "inputs": inputs_dict, - "abundances": abundances, - "T_e": 1 * u.MK, - "n": n_array, - "time_input": time_array, - "time_start": 0 * u.s, - "time_max": 800 * u.s, - "dt": 100 * u.s, - "max_steps": 2, - "adapt_dt": False, - "verbose": True, - }), - ("n_e constant", { - "inputs": inputs_dict, - "abundances": abundances, - "T_e": T_e_array, - "n": 1e9 * u.cm ** -3, - "time_input": time_array, - "time_start": 0 * u.s, - "time_max": 800 * u.s, - "max_steps": 2, - "adapt_dt": False, - "dt": 100 * u.s, - "verbose": True, - }), - ("T_e function", { - "inputs": inputs_dict, - "abundances": abundances, - "T_e": lambda time: 1e4 * (1 + time / u.s) * u.K, - "n": 1e15 * u.cm ** -3, - "time_max": 800 * u.s, - "max_steps": 2, - "dt": 100 * u.s, - "adapt_dt": False, - "verbose": True, - }), - ("n function", { - "inputs": inputs_dict, - "abundances": abundances, - "T_e": 6e4 * u.K, - "n": lambda time: 1e9 * (1 + time / u.s) * u.cm ** -3, - "time_start": 0 * u.s, - "time_max": 800 * u.s, - "adapt_dt": False, - "dt": 200 * u.s, - "verbose": True, - "max_steps": 4, - }), - ("equil test cool", { - "inputs": ["H", "He", "N"], - "abundances": {"H": 1, "He": 0.1, "C": 1e-4, "N": 1e-4, "O": 1e-4, "Fe": 1e-4}, - "T_e": 10001.0 * u.K, - "n": 1e13 * u.cm ** -3, - "time_max": 2e6 * u.s, - "tol": 1e-9, - "adapt_dt": False, - "dt": 5e5 * u.s, - "max_steps": 4, - "verbose": True, - }), - ("equil test hot", { - "inputs": ["H", "He", "C"], - "abundances": { - "H": 1, - "He": 0.1, - "C": 1e-4, - "N": 1e-4, - "O": 1e-4, - "Fe": 1e-4, - "S": 2e-6, - }, - "T_e": 7e6 * u.K, - "n": 1e9 * u.cm ** -3, - "time_max": 1e8 * u.s, - "dt": 5e7 * u.s, - "max_steps": 3, - "adapt_dt": False, - "verbose": True, - }), - ("equil test start far out of equil", { - "inputs": { - "H": [0.99, 0.01], - "He": [0.5, 0.0, 0.5], - "O": [0.2, 0, 0.2, 0, 0.2, 0, 0.2, 0, 0.2], - }, - "abundances": {"H": 1, "He": 0.1, "O": 1e-4}, - "T_e": 3e6 * u.K, - "n": 1e9 * u.cm ** -3, - "dt": 1e6 * u.s, - "time_start": 0 * u.s, - "time_max": 1e6 * u.s, - "adapt_dt": False, - "verbose": True, - "max_steps": 2, - }), - ("adapt dt", { - "inputs": ["H", "He"], - "abundances": {"H": 1, "He": 0.1}, - "T_e": lambda t: u.K * (1e6 + 1.3e4 * np.sin(t.value)), - "n": 1e10 * u.cm ** -3, - "max_steps": 300, - "time_start": 0 * u.s, - "time_max": 2 * np.pi * u.s, - "adapt_dt": True, - }), -]) +@pytest.fixture( + scope="module", + params=[ + ( + "basic", + { + "inputs": inputs_dict, + "abundances": abundances, + "T_e": T_e_array, + "n": n_array, + "time_input": time_array, + "time_start": 0 * u.s, + "time_max": 800 * u.s, + "max_steps": 1, + "dt": 800 * u.s, + "adapt_dt": False, + "verbose": True, + }, + ), + ( + "T_e constant", + { + "inputs": inputs_dict, + "abundances": abundances, + "T_e": 1 * u.MK, + "n": n_array, + "time_input": time_array, + "time_start": 0 * u.s, + "time_max": 800 * u.s, + "dt": 100 * u.s, + "max_steps": 2, + "adapt_dt": False, + "verbose": True, + }, + ), + ( + "n_e constant", + { + "inputs": inputs_dict, + "abundances": abundances, + "T_e": T_e_array, + "n": 1e9 * u.cm ** -3, + "time_input": time_array, + "time_start": 0 * u.s, + "time_max": 800 * u.s, + "max_steps": 2, + "adapt_dt": False, + "dt": 100 * u.s, + "verbose": True, + }, + ), + ( + "T_e function", + { + "inputs": inputs_dict, + "abundances": abundances, + "T_e": lambda time: 1e4 * (1 + time / u.s) * u.K, + "n": 1e15 * u.cm ** -3, + "time_max": 800 * u.s, + "max_steps": 2, + "dt": 100 * u.s, + "adapt_dt": False, + "verbose": True, + }, + ), + ( + "n function", + { + "inputs": inputs_dict, + "abundances": abundances, + "T_e": 6e4 * u.K, + "n": lambda time: 1e9 * (1 + time / u.s) * u.cm ** -3, + "time_start": 0 * u.s, + "time_max": 800 * u.s, + "adapt_dt": False, + "dt": 200 * u.s, + "verbose": True, + "max_steps": 4, + }, + ), + ( + "equil test cool", + { + "inputs": ["H", "He", "N"], + "abundances": { + "H": 1, + "He": 0.1, + "C": 1e-4, + "N": 1e-4, + "O": 1e-4, + "Fe": 1e-4, + }, + "T_e": 10001.0 * u.K, + "n": 1e13 * u.cm ** -3, + "time_max": 2e6 * u.s, + "tol": 1e-9, + "adapt_dt": False, + "dt": 5e5 * u.s, + "max_steps": 4, + "verbose": True, + }, + ), + ( + "equil test hot", + { + "inputs": ["H", "He", "C"], + "abundances": { + "H": 1, + "He": 0.1, + "C": 1e-4, + "N": 1e-4, + "O": 1e-4, + "Fe": 1e-4, + "S": 2e-6, + }, + "T_e": 7e6 * u.K, + "n": 1e9 * u.cm ** -3, + "time_max": 1e8 * u.s, + "dt": 5e7 * u.s, + "max_steps": 3, + "adapt_dt": False, + "verbose": True, + }, + ), + ( + "equil test start far out of equil", + { + "inputs": { + "H": [0.99, 0.01], + "He": [0.5, 0.0, 0.5], + "O": [0.2, 0, 0.2, 0, 0.2, 0, 0.2, 0, 0.2], + }, + "abundances": {"H": 1, "He": 0.1, "O": 1e-4}, + "T_e": 3e6 * u.K, + "n": 1e9 * u.cm ** -3, + "dt": 1e6 * u.s, + "time_start": 0 * u.s, + "time_max": 1e6 * u.s, + "adapt_dt": False, + "verbose": True, + "max_steps": 2, + }, + ), + ( + "adapt dt", + { + "inputs": ["H", "He"], + "abundances": {"H": 1, "He": 0.1}, + "T_e": lambda t: u.K * (1e6 + 1.3e4 * np.sin(t.value)), + "n": 1e10 * u.cm ** -3, + "max_steps": 300, + "time_start": 0 * u.s, + "time_max": 2 * np.pi * u.s, + "adapt_dt": True, + }, + ), + ], +) def name_inputs_instance(request): test_name, dictionary = request.param return test_name, dictionary, NEI(**dictionary) @@ -154,27 +191,32 @@ def test_time_start(name_inputs_instance): else: assert instance.time_start == 0 * u.s + def test_time_max(name_inputs_instance): test_name, inputs, instance = name_inputs_instance if "time_max" in inputs.keys(): assert inputs["time_max"] == instance.time_max + def test_initial_type(name_inputs_instance): test_name, inputs, instance = name_inputs_instance assert isinstance(instance.initial, IonizationStates) + def test_n_input(name_inputs_instance): test_name, inputs, instance = name_inputs_instance actual = instance.n_input expected = inputs["n"] assert np.all(expected == actual) + def test_T_e_input(name_inputs_instance): test_name, inputs, instance = name_inputs_instance actual = instance.T_e_input expected = inputs["T_e"] assert np.all(expected == actual) + def test_electron_temperature(name_inputs_instance): test_name, inputs, instance = name_inputs_instance T_e_input = instance.T_e_input @@ -188,27 +230,26 @@ def test_electron_temperature(name_inputs_instance): assert np.isclose(T_e.value, T_e_func.value) if callable(T_e_input): - assert instance.T_e_input( + assert instance.T_e_input(instance.time_start) == instance.electron_temperature( instance.time_start - ) == instance.electron_temperature(instance.time_start) + ) + def test_initial_ionfracs(name_inputs_instance): test_name, inputs, instance = name_inputs_instance if not isinstance(inputs["inputs"], dict): pytest.skip("Test irrelevant") - original_inputs = inputs["inputs"] original_elements = original_inputs.keys() for element in original_elements: assert np.allclose( original_inputs[element], - instance.initial.ionic_fractions[ - particle_symbol(element) - ], + instance.initial.ionic_fractions[particle_symbol(element)], ) + def test_simulate(name_inputs_instance): test_name, inputs, instance = name_inputs_instance instance.simulate() @@ -222,6 +263,7 @@ def test_simulate(name_inputs_instance): results.ionic_fractions[elem][-1, :], final.ionic_fractions[elem] ) + def test_simulation_end(name_inputs_instance): test_name, inputs, instance = name_inputs_instance time = instance.results.time @@ -243,6 +285,7 @@ def test_simulation_end(name_inputs_instance): print(f"time_max = {instance.time_max}") raise Exception("Problem with end time.") + def test_equilibration(name_inputs_instance): test_name, inputs, instance = name_inputs_instance if "equil" in test_name: @@ -259,6 +302,7 @@ def test_equilibration(name_inputs_instance): equil_dict[element], instance.results.ionic_fractions[element][-1, :] ) + def test_initial_results(name_inputs_instance): test_name, inputs, instance = name_inputs_instance initial = instance.initial From cb6378c1c3550cad5c9b5f556d582c0b7ffed527 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Wed, 3 Jun 2020 15:20:27 -0400 Subject: [PATCH 48/49] Improve conditional statements for T_e & T_e_index --- plasmapy_nei/eigen/eigenclass.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/plasmapy_nei/eigen/eigenclass.py b/plasmapy_nei/eigen/eigenclass.py index cda3a89..3b47b3f 100644 --- a/plasmapy_nei/eigen/eigenclass.py +++ b/plasmapy_nei/eigen/eigenclass.py @@ -362,9 +362,9 @@ def eigenvectors(self, T_e: u.K = None, T_e_index: u.K = None): # TODO: add discussion of what the indices of the returned array represent - if T_e_index: + if T_e_index is not None: return self._eigenvectors[T_e_index, :, :] - elif T_e: + elif T_e is not None: T_e_index = self._get_temperature_index(T_e) return self._eigenvectors[T_e_index, :, :] elif self.temperature: @@ -386,9 +386,9 @@ def eigenvector_inverses(self, T_e=None, T_e_index=None): The index of the electron temperature array corresponding to the desired temperature. """ - if T_e_index: + if T_e_index is not None: return self._eigenvector_inverses[T_e_index, :, :] - elif T_e: + elif T_e is not None: T_e_index = self._get_temperature_index(T_e) return self._eigenvector_inverses[T_e_index, :, :] elif self.temperature: @@ -410,9 +410,9 @@ def equilibrium_state(self, T_e=None, T_e_index=None): The index of the electron temperature array corresponding to the desired temperature. """ - if T_e_index: + if T_e_index is not None: return self._equilibrium_states[T_e_index, :] - elif T_e: + elif T_e is not None: T_e_index = self._get_temperature_index(T_e) return self._equilibrium_states[T_e_index, :] elif self.temperature: From 293d31939af1aa6a8676f21db53ae42d86d532a4 Mon Sep 17 00:00:00 2001 From: Nick Murphy Date: Wed, 3 Jun 2020 17:55:38 -0400 Subject: [PATCH 49/49] Update plasmapy_nei/nei/nei.py Co-authored-by: Erik Everson --- plasmapy_nei/nei/nei.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plasmapy_nei/nei/nei.py b/plasmapy_nei/nei/nei.py index 6b7e152..35103a8 100644 --- a/plasmapy_nei/nei/nei.py +++ b/plasmapy_nei/nei/nei.py @@ -848,7 +848,7 @@ def max_steps(self, n: int): self._max_steps = n else: raise TypeError( - "max_steps must be an integer with 0 < max_steps <= " "1000000" + "max_steps must be an integer with 0 < max_steps <= 1000000" ) @property