unittest for omas_physics

omas/omas_sample.py

@@ -220,7 +220,7 @@ def equilibrium(ods, time_index=0, include_profiles=False, include_phi=False, in
 
 
 @add_to_ODS
-def profiles(ods, time_index=0, nx=11, add_junk_ion=False):
+def profiles(ods, time_index=0, nx=11, add_junk_ion=False, include_pressure=True):
     """
     Add made up sample profiles to an ODS. Although made up, the profiles satisfy quasi-neutrality and should very
     roughly resemble something similar to a real plasma a little bit.
@@ -235,6 +235,9 @@ def profiles(ods, time_index=0, nx=11, add_junk_ion=False):
     :param add_junk_ion: bool
         Flag for adding a junk ion for testing how well functions tolerate problems. This will be missing labels, etc.
 
+    :param include_pressure: bool
+        Include pressure profiles when temperature and density are added
+
     :return: ODS instance with profiles added.
         Since the original is modified, it is not necessary to catch the return, but it may be convenient to do so in
         some contexts. If you do not want the original to be modified, deepcopy it first.
@@ -247,8 +250,9 @@ def profiles(ods, time_index=0, nx=11, add_junk_ion=False):
     prof1d['electrons.density'] = 1e19 * (6.5 - 1.9*x - 4.5*x**7)  # m^-3
     prof1d['electrons.density_thermal'] = prof1d['electrons.density']
     prof1d['electrons.temperature'] = 1000 * (4 - 3*x - 0.9*x**9)  # eV
-    prof1d['electrons.pressure'] = prof1d['electrons.density'] * prof1d['electrons.temperature'] * ee
-    prof1d['electrons.pressure_thermal'] = prof1d['electrons.density_thermal'] * prof1d['electrons.temperature'] * ee
+    if include_pressure:
+        prof1d['electrons.pressure'] = prof1d['electrons.density'] * prof1d['electrons.temperature'] * ee
+        prof1d['electrons.pressure_thermal'] = prof1d['electrons.density_thermal'] * prof1d['electrons.temperature'] * ee
 
     prof1d['ion.0.label'] = 'D+'
     prof1d['ion.0.element.0.z_n'] = 1.0
@@ -258,8 +262,9 @@ def profiles(ods, time_index=0, nx=11, add_junk_ion=False):
     prof1d['ion.0.density_fast'] = prof1d['ion.0.density'].max() * 0.32 * numpy.exp(-(x**2)/0.3**2/2.)
     prof1d['ion.0.density_thermal'] = prof1d['ion.0.density'] - prof1d['ion.0.density_fast']
     prof1d['ion.0.temperature'] = prof1d['electrons.temperature'] * 1.1
-    prof1d['ion.0.pressure'] = prof1d['ion.0.temperature'] * prof1d['ion.0.density'] * ee
-    prof1d['ion.0.pressure_thermal'] = prof1d['ion.0.temperature'] * prof1d['ion.0.density_thermal'] * ee
+    if include_pressure:
+        prof1d['ion.0.pressure'] = prof1d['ion.0.temperature'] * prof1d['ion.0.density'] * ee
+        prof1d['ion.0.pressure_thermal'] = prof1d['ion.0.temperature'] * prof1d['ion.0.density_thermal'] * ee
 
     prof1d['ion.1.label'] = 'C+6'
     prof1d['ion.1.element.0.z_n'] = 6.0
@@ -268,17 +273,19 @@ def profiles(ods, time_index=0, nx=11, add_junk_ion=False):
     prof1d['ion.1.density'] = (prof1d['electrons.density'] - prof1d['ion.0.density']) / prof1d['ion.1.element.0.z_n']
     prof1d['ion.1.density_thermal'] = prof1d['ion.1.density']
     prof1d['ion.1.temperature'] = prof1d['ion.0.temperature']*0.98
-    prof1d['ion.1.pressure'] = prof1d['ion.1.temperature'] * prof1d['ion.1.density'] * ee
-    prof1d['ion.1.pressure_thermal'] = prof1d['ion.1.temperature'] * prof1d['ion.1.density_thermal'] * ee
+    if include_pressure:
+        prof1d['ion.1.pressure'] = prof1d['ion.1.temperature'] * prof1d['ion.1.density'] * ee
+        prof1d['ion.1.pressure_thermal'] = prof1d['ion.1.temperature'] * prof1d['ion.1.density_thermal'] * ee
 
     if add_junk_ion:
         junki = prof1d['ion.2']
         junki['density'] = x*0
         junki['temperature'] = x*0
-        junki['pressure'] = x*0
 
         junkn = prof1d['neutral.0']
-        junkn['pressure'] = x*0
+        if include_pressure:
+            junki['pressure'] = x*0
+            junkn['pressure'] = x*0
 
     return ods
 

omas/omas_utils.py

@@ -34,7 +34,7 @@ def different_ods(ods1, ods2):
         elif type(ods1[k]) != type(ods2[k]):
             return 'DIFF: `%s` differ in type (%s,%s)' % (k, type(ods1[k]), type(ods2[k]))
         elif numpy.atleast_1d(is_uncertain(ods1[k])).any() or numpy.atleast_1d(is_uncertain(ods2[k])).any():
-            if not numpy.allclose(nominal_values(ods1[k]), nominal_values(ods2[k])) and not numpy.allclose(std_devs(ods1[k]), std_devs(ods2[k])):
+            if not numpy.allclose(nominal_values(ods1[k]), nominal_values(ods2[k])) or not numpy.allclose(std_devs(ods1[k]), std_devs(ods2[k])):
                 return 'DIFF: `%s` differ in value' % k
         else:
             if not numpy.allclose(ods1[k], ods2[k]):

omas/tests/test_omas_physics.py

@@ -0,0 +1,93 @@
+#!/usr/bin/env python
+# # -*- coding: utf-8 -*-
+
+"""
+Test script for omas/omas_physics.py
+python -m unittest test_omas_physics (from omas/omas/tests)
+python -m unittest discover omas (from omas top level; runs all tests)
+"""
+
+# Basic imports
+from __future__ import print_function, division, unicode_literals
+import os
+import unittest
+import numpy
+import warnings
+import copy
+
+# Plot imports
+import matplotlib as mpl
+from matplotlib import pyplot as plt
+
+# OMAS imports
+from omas import *
+from omas.omas_utils import *
+from omas.omas_physics import *
+
+
+class TestOmasPhysics(unittest.TestCase):
+    """
+    Test suite for omas_physics.py
+    """
+
+    # Flags to edit while testing
+    verbose = False  # Spammy, but occasionally useful for debugging a weird problem
+
+    # Sample data for use in tests
+    ods = ODS()
+
+    # Utilities for this test
+    def printv(self, *arg):
+        """Utility for tests to use"""
+        if self.verbose:
+            print(*arg)
+
+    def setUp(self):
+        test_id = self.id()
+        test_name = '.'.join(test_id.split('.')[-2:])
+        self.printv('{}...'.format(test_name))
+
+    def tearDown(self):
+        test_name = '.'.join(self.id().split('.')[-2:])
+        self.printv('    {} done.'.format(test_name))
+
+    @unittest.expectedFailure  # TODO: fix core_profiles_pressures or revise this test
+    def test_core_profiles_pressures(self):
+        ods2 = copy.deepcopy(self.ods)
+        ods2.sample_profiles(include_pressure=False)
+        ods3 = copy.deepcopy(self.ods)
+        updated_ods = core_profiles_pressures(ods2, update=False)
+        updated_ods3 = core_profiles_pressures(ods3, update=True)
+
+        assert updated_ods3 == ods3
+        assert updated_ods != ods2
+        assert updated_ods3 == updated_ods
+
+    def test_define_cocos(self):
+        cocos_none = define_cocos(None)
+        cocos1 = define_cocos(1)
+        cocos2 = define_cocos(2)
+        cocos3 = define_cocos(3)
+        cocos4 = define_cocos(4)
+        cocos5 = define_cocos(5)
+        cocos6 = define_cocos(6)
+        cocos7 = define_cocos(7)
+        cocos8 = define_cocos(8)
+        cocos11 = define_cocos(11)
+        for cocos in [cocos_none, cocos1, cocos2, cocos5, cocos6, cocos11]:
+            assert cocos['sigma_Bp'] == 1
+        for cocos in [cocos3, cocos4, cocos7, cocos8]:
+            assert cocos['sigma_Bp'] == -1
+
+    def test_cocos_transform(self):
+        assert cocos_transform(None, None)['TOR'] == 1
+        assert cocos_transform(1, 3)['POL'] == -1
+        for cocos_ind in range(1, 9):
+            assert cocos_transform(cocos_ind, cocos_ind + 10)['invPSI'] != 1
+            for cocos_add in range(2):
+                for thing in ['BT', 'TOR', 'POL', 'Q']:
+                    assert cocos_transform(cocos_ind+cocos_add*10, cocos_ind+cocos_add*10)[thing] == 1
+
+
+if __name__ == '__main__':
+    unittest.main()

omas/tests/test_omas_plot.py

@@ -3,7 +3,8 @@
 
 """
 Test script for omas/omas_plot.py
-python -m unittest test_omas_plot
+python -m unittest test_omas_plot (from omas/omas/tests)
+python -m unittest discover omas (from omas top level; runs all tests)
 """
 
 # Basic imports

omas/tests/test_omas_utils.py

@@ -0,0 +1,156 @@
+#!/usr/bin/env python
+# # -*- coding: utf-8 -*-
+
+"""
+Test script for omas/omas_utils.py
+python -m unittest test_omas_utils (from omas/omas/tests)
+python -m unittest discover omas (from omas top level; runs all tests)
+"""
+
+# Basic imports
+from __future__ import print_function, division, unicode_literals
+import os
+import unittest
+import numpy
+import warnings
+import copy
+
+# Plot imports
+import matplotlib as mpl
+from matplotlib import pyplot as plt
+
+# OMAS imports
+from omas import *
+from omas.omas_utils import *
+
+
+class TestOmasUtils(unittest.TestCase):
+    """
+    Test suite for omas_utils.py
+    """
+
+    # Flags to edit while testing
+    verbose = False  # Spammy, but occasionally useful for debugging a weird problem
+
+    # Sample data for use in tests
+    ods = ODS()
+    specific_test_version = '3.18.0'
+
+    # Utilities for this test
+    def printv(self, *arg):
+        """Utility for tests to use"""
+        if self.verbose:
+            print(*arg)
+
+    def setUp(self):
+        test_id = self.id()
+        test_name = '.'.join(test_id.split('.')[-2:])
+        self.printv('{}...'.format(test_name))
+
+    def tearDown(self):
+        test_name = '.'.join(self.id().split('.')[-2:])
+        self.printv('    {} done.'.format(test_name))
+
+    def test_different_ods(self):
+        assert different_ods(self.ods, self.ods) is False
+        ods2 = ODS()
+        ods2.sample_equilibrium()
+        diff_eq = different_ods(self.ods, ods2)
+        self.printv('  diff_eq = {}'.format(diff_eq))
+        assert isinstance(diff_eq, basestring)
+        assert ('equilibrium' in diff_eq) or ('wall' in diff_eq)
+        ods3 = copy.deepcopy(ods2)
+        assert different_ods(ods2, ods3) is False
+        ods3.sample_profiles()
+        diff_prof = different_ods(ods2, ods3)
+        self.printv('  diff_prof = {}'.format(diff_prof))
+        assert isinstance(diff_prof, basestring)
+        assert isinstance(different_ods(ods3, ods2), basestring)
+        assert 'core_profiles' in diff_prof
+        ods2.sample_profiles(include_pressure=False)
+        diff_prof2 = different_ods(ods3, ods2)
+        self.printv('  diff_prof2 = {}'.format(diff_prof2))
+        assert isinstance(diff_prof2, basestring)
+        assert 'core_profiles' in diff_prof2
+        ods2.sample_profiles()
+        ods2['core_profiles.profiles_1d.0.electrons.density'][0] = 1.5212
+        diff_prof3 = different_ods(ods2, ods3)
+        self.printv('  diff_prof3 = {}'.format(diff_prof3))
+        assert isinstance(diff_prof3, basestring)
+        assert 'value' in diff_prof3
+        ods2.sample_profiles()
+        ods2['core_profiles.profiles_1d.0.ion.0.element.0.a'] = 9.
+        diff_prof4 = different_ods(ods2, ods3)
+        self.printv('  diff_prof4 = {}'.format(diff_prof4))
+        assert isinstance(diff_prof4, basestring)
+        assert 'value' in diff_prof4
+        ods2.sample_profiles()
+        ods2['core_profiles.code.name'] = 'fake name 1'
+        ods3['core_profiles.code.name'] = 'fake name 2'
+        diff_prof5 = different_ods(ods2, ods3)
+        self.printv('  diff_prof5 = {}'.format(diff_prof5))
+        assert isinstance(diff_prof5, basestring)
+        assert 'name' in diff_prof5
+        ods3['core_profiles.code.name'] = numpy.array([2, 3, 4])
+        assert isinstance(different_ods(ods2, ods3), basestring)
+        ods2['core_profiles.code.name'] = uarray(numpy.array([2, 3, 4]), numpy.array([1, 1, 1]))
+        assert isinstance(different_ods(ods2, ods3), basestring)
+
+    def test_printe(self):
+        printe('printe_test,', end='')
+        printw('printw_test', end='')
+
+    def test_is_numeric(self):
+        assert is_numeric(5) is True
+        assert is_numeric(numpy.array([5])) is True
+        assert is_numeric('blah') is False
+        assert is_numeric({'blah': 'blah'}) is False
+        assert is_numeric([]) is False
+        assert is_numeric(None) is False
+
+    def test_remove_parentheses(self):
+        assert remove_parentheses('zoom(blah)what', replace_with='|') == 'zoom|what'
+
+    def test_closest_index(self):
+        # Basic tests
+        assert closest_index([1, 2, 3, 4], 3) == 2  # Basic test
+        assert closest_index([1, 2, 3], 1) == 0  # Special: The first element is the one sought
+        assert closest_index(numpy.array([1, 2, 3, 4]), 4) == 3  # Special: The last element is the one sought
+        assert closest_index([1, 2, 2, 3], 2) == 1  # Special: duplicated value: pick first instance
+        assert closest_index([1, 2, 3, 4], 2.2) == 1  # Make sure it works for numbers in between
+        assert closest_index([1, 2, 3, 4], 2.7) == 2
+        # Exception handling and coping with problems
+        self.assertRaises(TypeError, closest_index, 5, 5)  # First arg is not a list --> TypeError
+        self.assertRaises(TypeError, closest_index, [1, 2, 3], 'string_not_number')
+        assert closest_index([1, 2, 3], [3]) == 2  # Should use first element of second arg if it's not a scalar
+        self.assertRaises(TypeError, closest_index, [1, 2, 3], [3, 2, 1])  # Can't call w/ list as 2nd arg unless len=1
+
+    def test_list_structures(self):  # Also tests dict_structures
+        struct_list = list_structures(default_imas_version)
+        struct_list2 = list_structures(self.specific_test_version)
+        assert isinstance(struct_list, list)
+        assert isinstance(struct_list2, list)
+        assert isinstance(struct_list[0], basestring)
+        assert 'pf_active' in struct_list2
+        struct_dict = dict_structures(default_imas_version)
+        struct_dict2 = dict_structures(self.specific_test_version)
+        assert isinstance(struct_dict, dict)
+        assert isinstance(struct_dict2, dict)
+        assert 'pf_active' in struct_dict2.keys()
+        assert all([item in struct_dict.keys() for item in struct_list])
+        assert all([item in struct_dict2.keys() for item in struct_list2])
+
+    def test_omas_info(self):
+        get_list = ['pf_active', 'thomson_scattering', 'charge_exchange']
+        ods_info_pfa = omas_info(get_list[0])
+        ods_info_pfa2 = omas_info(get_list[0], self.specific_test_version)
+        assert get_list[0] in ods_info_pfa.keys()
+        assert get_list[0] in ods_info_pfa2.keys()
+        if get_list[0] == 'pf_active':
+            assert isinstance(ods_info_pfa['pf_active.circuit.0.connections.documentation'], basestring)
+        ods_info_list = omas_info(get_list)
+        assert all([item in ods_info_list for item in get_list])
+
+
+if __name__ == '__main__':
+    unittest.main()