From dedbbc8822c4d1c500bb9cb25747b471ac0a9a64 Mon Sep 17 00:00:00 2001 From: wiltbemj Date: Fri, 21 Feb 2025 12:29:46 -0700 Subject: [PATCH] Adds tests for RCM --- tests/rcm/lambdautils/test_AlamData.py | 80 +++++++++++++++++++++++ tests/rcm/lambdautils/test_AlamParams.py | 44 +++++++++++++ tests/rcm/lambdautils/test_fileIO.py | 82 ++++++++++++++++++++++++ tests/rcm/lambdautils/test_genAlam.py | 71 ++++++++++++++++++++ tests/rcm/test_rcminit.py | 67 +++++++++++++++++++ tests/rcm/test_rcmutils.py | 73 +++++++++++++++++++++ tests/rcm/wmutils/test_genWM.py | 68 ++++++++++++++++++++ tests/rcm/wmutils/test_wmData.py | 30 +++++++++ 8 files changed, 515 insertions(+) create mode 100644 tests/rcm/lambdautils/test_AlamData.py create mode 100644 tests/rcm/lambdautils/test_AlamParams.py create mode 100644 tests/rcm/lambdautils/test_fileIO.py create mode 100644 tests/rcm/lambdautils/test_genAlam.py create mode 100644 tests/rcm/test_rcminit.py create mode 100644 tests/rcm/test_rcmutils.py create mode 100644 tests/rcm/wmutils/test_genWM.py create mode 100644 tests/rcm/wmutils/test_wmData.py diff --git a/tests/rcm/lambdautils/test_AlamData.py b/tests/rcm/lambdautils/test_AlamData.py new file mode 100644 index 0000000..1e5366e --- /dev/null +++ b/tests/rcm/lambdautils/test_AlamData.py @@ -0,0 +1,80 @@ +import pytest +from kaipy.rcm.lambdautils.AlamData import Species, AlamData +from kaipy.rcm.lambdautils.AlamParams import AlamParams, SpecParams +import kaipy.rcm.lambdautils.DistTypes as dT + + +def setup_AlamParams(): + spec_params = [ + SpecParams(n=10, amin=0.1, amax=1.0, distType=dT.DistType(), flav=1), + SpecParams(n=20, amin=0.2, amax=2.0, distType=dT.DistType(), flav=2) + ] + params = AlamParams( + doUsePsphere=True, + specParams=spec_params, + emine=0.01, + eminp=0.02, + emaxe=10.0, + emaxp=20.0, + L_kt=5.0 + ) + return params + +def test_species_initialization(): + params = setup_AlamParams() + species = Species( + n=3, + alams=[1.0, 2.0, 3.0], + amins=[0.5, 1.5, 2.5], + amaxs=[1.5, 2.5, 3.5], + flav=1, + fudge=0.1, + params=params, + name="TestSpecies" + ) + + + assert species.n == 3 + assert species.alams == [1.0, 2.0, 3.0] + assert species.amins == [0.5, 1.5, 2.5] + assert species.amaxs == [1.5, 2.5, 3.5] + assert species.flav == 1 + assert species.fudge == 0.1 + assert isinstance(species.params, AlamParams) + assert species.name == "TestSpecies" + +def test_alamdata_initialization(): + params = setup_AlamParams() + species_list = [ + Species( + n=3, + alams=[1.0, 2.0, 3.0], + amins=[0.5, 1.5, 2.5], + amaxs=[1.5, 2.5, 3.5], + flav=1, + fudge=0.1, + name="TestSpecies1" + ), + Species( + n=2, + alams=[4.0, 5.0], + amins=[3.5, 4.5], + amaxs=[4.5, 5.5], + flav=2, + fudge=0.2, + name="TestSpecies2" + ) + ] + + alam_data = AlamData( + doUsePsphere=True, + specs=species_list, + params=params + ) + + assert alam_data.doUsePsphere is True + assert len(alam_data.specs) == 2 + assert isinstance(alam_data.specs[0], Species) + assert alam_data.specs[0].name == "TestSpecies1" + assert alam_data.specs[1].name == "TestSpecies2" + assert isinstance(alam_data.params, AlamParams) \ No newline at end of file diff --git a/tests/rcm/lambdautils/test_AlamParams.py b/tests/rcm/lambdautils/test_AlamParams.py new file mode 100644 index 0000000..4ed64df --- /dev/null +++ b/tests/rcm/lambdautils/test_AlamParams.py @@ -0,0 +1,44 @@ +import pytest +from kaipy.rcm.lambdautils.AlamParams import AlamParams, SpecParams + +import kaipy.rcm.lambdautils.DistTypes as dT + +def test_AlamParams_initialization(): + spec_params = [ + SpecParams(n=10, amin=0.1, amax=1.0, distType=dT.DistType(), flav=1), + SpecParams(n=20, amin=0.2, amax=2.0, distType=dT.DistType(), flav=2) + ] + params = AlamParams( + doUsePsphere=True, + specParams=spec_params, + emine=0.01, + eminp=0.02, + emaxe=10.0, + emaxp=20.0, + L_kt=5.0 + ) + + assert params.doUsePsphere is True + assert len(params.specParams) == 2 + assert params.emine == 0.01 + assert params.eminp == 0.02 + assert params.emaxe == 10.0 + assert params.emaxp == 20.0 + assert params.L_kt == 5.0 + +def test_AlamParams_default_values(): + spec_params = [ + SpecParams(n=10, amin=0.1, amax=1.0, distType=dT.DistType(), flav=1) + ] + params = AlamParams( + doUsePsphere=False, + specParams=spec_params + ) + + assert params.doUsePsphere is False + assert len(params.specParams) == 1 + assert params.emine is None + assert params.eminp is None + assert params.emaxe is None + assert params.emaxp is None + assert params.L_kt is None diff --git a/tests/rcm/lambdautils/test_fileIO.py b/tests/rcm/lambdautils/test_fileIO.py new file mode 100644 index 0000000..8c15393 --- /dev/null +++ b/tests/rcm/lambdautils/test_fileIO.py @@ -0,0 +1,82 @@ +import pytest +import numpy as np +import h5py +import os + +import kaipy.rcm.lambdautils.AlamParams as aP +from kaipy.rcm.lambdautils.fileIO import saveRCMConfig, saveParams, loadParams, saveData +from kaipy.rcm.lambdautils.AlamParams import AlamParams, SpecParams +from kaipy.rcm.lambdautils.AlamData import Species, AlamData +import kaipy.rcm.lambdautils.DistTypes as dT + +def setup_AlamParams(): + spec_params = [ + SpecParams(n=10, amin=0.1, amax=1.0, distType=dT.DistType(), flav=1), + SpecParams(n=20, amin=0.2, amax=2.0, distType=dT.DistType(), flav=2) + ] + params = AlamParams( + doUsePsphere=True, + specParams=spec_params, + emine=0.01, + eminp=0.02, + emaxe=10.0, + emaxp=20.0, + L_kt=5.0 + ) + return params + +def setup_AlamData(): + params = setup_AlamParams() + species_list = [ + Species( + n=3, + alams=[1.0, 2.0, 3.0], + amins=[0.5, 1.5, 2.5], + amaxs=[1.5, 2.5, 3.5], + flav=1, + fudge=0.1, + name="TestSpecies1" + ), + Species( + n=2, + alams=[4.0, 5.0], + amins=[3.5, 4.5], + amaxs=[4.5, 5.5], + flav=2, + fudge=0.2, + name="TestSpecies2" + ) + ] + alamData = AlamData( + doUsePsphere=True, + specs=species_list, + params=params + ) + return alamData + +def test_saveRCMConfig(tmp_path): + alamData = setup_AlamData() + fname = tmp_path / 'rcmconfig.h5' + saveRCMConfig(alamData, params=None, fname=fname) + assert os.path.exists(fname) + +def test_saveParams(tmp_path): + f5 = h5py.File(tmp_path / 'testfile.h5', 'w') + alamParams = setup_AlamParams() + saveParams(f5, alamParams) + assert 'AlamParams' in f5.attrs + f5.close() + +def test_loadParams(tmp_path): + f5 = h5py.File(tmp_path / 'testfile.h5', 'w') + alamParams = setup_AlamParams() + saveParams(f5, alamParams) + f5.close() + + f5 = h5py.File(tmp_path / 'testfile.h5', 'r') + loaded_params = loadParams(f5) + print(f'loaded_params: {loaded_params}') + print(f'alamParams: {alamParams}') + assert loaded_params['doUsePsphere'] == alamParams.doUsePsphere + + f5.close() diff --git a/tests/rcm/lambdautils/test_genAlam.py b/tests/rcm/lambdautils/test_genAlam.py new file mode 100644 index 0000000..c801aff --- /dev/null +++ b/tests/rcm/lambdautils/test_genAlam.py @@ -0,0 +1,71 @@ +import pytest +import numpy as np +from kaipy.rcm.lambdautils.genAlam import getAlamMinMax, genSpeciesFromParams, genAlamDataFromParams, genPsphereSpecies +from kaipy.rcm.lambdautils.AlamParams import SpecParams, AlamParams +from kaipy.rcm.lambdautils.AlamData import Species, AlamData +import kaipy.rcm.lambdautils.DistTypes as dT + +def setup_AlamParams(): + wolfP1 = 3 + wolfP2 = 1 + spec_params = [ + SpecParams(n=10, amin=0.1, amax=1.0, distType=dT.DT_Wolf(p1=wolfP1,p2=wolfP2), flav=1), + SpecParams(n=20, amin=0.2, amax=2.0, distType=dT.DT_Wolf(p1=wolfP1,p2=wolfP2), flav=2) + ] + params = AlamParams( + doUsePsphere=True, + specParams=spec_params, + emine=0.01, + eminp=0.02, + emaxe=10.0, + emaxp=20.0, + L_kt=5.0 + ) + return params + +def test_getAlamMinMax(): + alams = [1, 2, 3, 4] + amin, amax = getAlamMinMax(alams) + assert amin == [0, 1.5, 2.5, 3.5] + assert amax == [1.5, 2.5, 3.5, 4.5] + +def test_genSpeciesFromParams(): + specParams = SpecParams(n=10, amin=0.1, amax=1.0, distType=dT.DistType(), flav=1) + specParams.genAlams = lambda: [1, 2, 3] + specParams.flav = 'test_flav' + specParams.fudge = 0.1 + specParams.name = 'test_name' + + species = genSpeciesFromParams(specParams) + + assert isinstance(species, Species) + assert species.n == 3 + assert species.alams == [1, 2, 3] + assert species.amins == [0, 1.5, 2.5] + assert species.amaxs == [1.5, 2.5, 3.5] + assert species.flav == 'test_flav' + assert species.fudge == 0.1 + assert species.name == 'test_name' + +def test_genAlamDataFromParams(): + params = setup_AlamParams() + + alamData = genAlamDataFromParams(params) + + assert alamData.doUsePsphere == True + assert len(alamData.specs) == 3 + assert isinstance(alamData.specs[0], Species) + assert alamData.specs[0].name == "Plasmasphere" + assert isinstance(alamData.params, AlamParams) + +def test_genPsphereSpecies(): + species = genPsphereSpecies() + + assert isinstance(species, Species) + assert species.n == 1 + assert species.alams == [0] + assert species.amins == [0] + assert species.amaxs == [0] + assert species.flav == 1 + assert species.fudge == 0 + assert species.name == 'Plasmasphere' \ No newline at end of file diff --git a/tests/rcm/test_rcminit.py b/tests/rcm/test_rcminit.py new file mode 100644 index 0000000..50c74da --- /dev/null +++ b/tests/rcm/test_rcminit.py @@ -0,0 +1,67 @@ +import pytest +import numpy as np +import os +from kaipy.rcm.rcminit import LoadLAS1 + +# filepath: /glade/u/home/wiltbemj/src/kaipy-private/kaipy/rcm/test_rcminit.py + +@pytest.fixture +def mock_las1_file(tmpdir): + data = """# Header line 1 +# Header line 2 +1.0 2.0 3.0 4.0 +5.0 6.0 7.0 8.0 +9.0 10.0 11.0 12.0 +""" + file_path = tmpdir.join("rcmlas1") + with open(file_path, 'w') as f: + f.write(data) + return str(file_path) + +def test_LoadLAS1(mock_las1_file, monkeypatch): + def mock_realpath(path): + return os.path.dirname(mock_las1_file) + + monkeypatch.setattr(os.path, 'realpath', mock_realpath) + + alamc, etac, ikflavc, fudgec = LoadLAS1(fIn="rcmlas1") + + assert np.array_equal(alamc, np.array([1.0, 5.0, 9.0])) + assert np.array_equal(etac, np.array([2.0, 6.0, 10.0])) + assert np.array_equal(ikflavc, np.array([3.0, 7.0, 11.0])) + assert np.array_equal(fudgec, np.array([4.0, 8.0, 12.0])) + +def test_LoadLAS1_file_not_found(monkeypatch): + def mock_realpath(path): + return "/non/existent/path" + + monkeypatch.setattr(os.path, 'realpath', mock_realpath) + + with pytest.raises(OSError): + LoadLAS1(fIn="nonexistentfile") + +def test_LoadLAS1_invalid_format(tmpdir, monkeypatch): + data = """# Header line 1 +# Header line 2 +1.0 2.0 +5.0 6.0 7.0 8.0 +""" + file_path = tmpdir.join("rcmlas1") + with open(file_path, 'w') as f: + f.write(data) + + def mock_realpath(path): + return os.path.dirname(str(file_path)) + + monkeypatch.setattr(os.path, 'realpath', mock_realpath) + + with pytest.raises(ValueError): + LoadLAS1(fIn="rcmlas1") + +def test_LoadLAS1_default(): + alamc, etac, ikflavc, fudgec = LoadLAS1() + assert len(alamc) == 90 + assert len(etac) == 90 + assert len(ikflavc) == 90 + assert len(fudgec) == 90 + diff --git a/tests/rcm/test_rcmutils.py b/tests/rcm/test_rcmutils.py new file mode 100644 index 0000000..ac536c2 --- /dev/null +++ b/tests/rcm/test_rcmutils.py @@ -0,0 +1,73 @@ +import pytest +import numpy as np +import h5py +from kaipy.rcm.rcmutils import getSpecieslambdata, getClosedRegionMask, getCumulPress, getValAtLoc_linterp, RCMSpeciesInfo, RCMInfo + + +def test_getSpecieslambdata(): + rcmS0 = { + 'alamc': np.array([0, 1, 2, 3, 4]) + } + result = getSpecieslambdata(rcmS0, species='ions') + assert result['kStart'] == 1 + assert result['kEnd'] == 5 + assert np.array_equal(result['ilamc'], np.array([1, 2, 3, 4])) + assert np.array_equal(result['ilami'], np.array([0, 1.5, 2.5, 3.5, 4.5])) + assert np.array_equal(result['lamscl'], np.diff(result['ilami']) * np.sqrt(result['ilamc'])) + +def test_getClosedRegionMask(): + s5 = { + 'rcmxmin': np.array([0.5, 1.5]), + 'rcmymin': np.array([0.5, 1.5]), + 'rcmvm': np.array([0, 1]) + } + mask = getClosedRegionMask(s5) + assert np.array_equal(mask, np.array([True, False])) + +def test_getCumulPress(): + ilamc = np.array([1, 2, 3]) + vm = np.array([[1, 2], [3, 4]]) + eetas = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]]) + pCumul = getCumulPress(ilamc, vm, eetas) + assert pCumul.shape == (3, 2, 2) + +def test_getValAtLoc_linterp(): + xData = [1, 2, 3, 4] + yData = [2, 4, 6, 8] + val = 2.5 + result = getValAtLoc_linterp(val, xData, yData, getAxis='y') + assert result == 5.0 + +def test_RCMSpeciesInfo_initialization(): + N = 5 + flav = 2 + kStart = 0 + kEnd = 5 + alamc = np.array([0, 1, 2, 3, 4]) + species_info = RCMSpeciesInfo(N, flav, kStart, kEnd, alamc) + + assert species_info.N == N + assert species_info.flav == flav + assert species_info.kStart == kStart + assert species_info.kEnd == kEnd + assert np.array_equal(species_info.alamc, alamc) + assert np.array_equal(species_info.alami, np.array([0, 0.5, 1.5, 2.5, 3.5, 4.5])) + +def test_RCMSpeciesInfo_single_cell(): + N = 1 + flav = 1 + kStart = 0 + kEnd = 1 + alamc = np.array([0]) + species_info = RCMSpeciesInfo(N, flav, kStart, kEnd, alamc) + + assert species_info.N == N + assert species_info.flav == flav + assert species_info.kStart == kStart + assert species_info.kEnd == kEnd + assert np.array_equal(species_info.alamc, alamc) + assert np.array_equal(species_info.alami, np.array([0, 0])) + + + + diff --git a/tests/rcm/wmutils/test_genWM.py b/tests/rcm/wmutils/test_genWM.py new file mode 100644 index 0000000..d5396af --- /dev/null +++ b/tests/rcm/wmutils/test_genWM.py @@ -0,0 +1,68 @@ +import pytest +import numpy as np +from kaipy.rcm.wmutils.genWM import genWM, genh5, readPoly, ChorusPoly, ReSample, genChorus +from kaipy.rcm.wmutils.wmData import wmParams + + +def test_ChorusPoly_basic(): + Li = np.array([3.0, 4.0, 5.0]) + Eki = np.array([0.001, 0.01, 0.1]) + polyArray = np.random.rand(7, 24, 10) # Mock polynomial coefficients + + result = ChorusPoly(Li, Eki, polyArray) + + assert result.shape == (7, 24, 3, 3) + assert np.all(result > 0) + +def test_ChorusPoly_zero_coefficients(): + Li = np.array([3.0, 4.0, 5.0]) + Eki = np.array([0.001, 0.01, 0.1]) + polyArray = np.zeros((7, 24, 10)) # All coefficients are zero + + result = ChorusPoly(Li, Eki, polyArray) + + assert result.shape == (7, 24, 3, 3) + assert np.all(result == 86400) # 10^0 * (60 * 60 * 24) = 86400 + +def test_ChorusPoly_specific_coefficients(): + Li = np.array([3.0, 4.0, 5.0]) + Eki = np.array([0.001, 0.01, 0.1]) + polyArray = np.zeros((7, 24, 10)) + polyArray[:, :, 0] = 1 # Set intercept to 1 + + result = ChorusPoly(Li, Eki, polyArray) + + assert result.shape == (7, 24, 3, 3) + assert np.all(result == 864000) # 10^1 * (60 * 60 * 24) = 864000 + +def test_ChorusPoly_varying_coefficients(): + Li = np.array([3.0, 4.0, 5.0]) + Eki = np.array([0.001, 0.01, 0.1]) + polyArray = np.zeros((7, 24, 10)) + polyArray[:, :, 1] = 1 # Set L coefficient to 1 + + result = ChorusPoly(Li, Eki, polyArray) + + assert result.shape == (7, 24, 3, 3) + assert np.all(result > 86400) # Result should be greater than 86400 due to L coefficient + +def test_ChorusPoly_large_values(): + Li = np.array([3.0, 4.0, 5.0]) + Eki = np.array([0.001, 0.01, 0.1]) + polyArray = np.ones((7, 24, 10)) * 100 # Large coefficients + + result = ChorusPoly(Li, Eki, polyArray) + + assert result.shape == (7, 24, 3, 3) + assert np.all(result > 0) + +def test_genWM_basic(): + params = wmParams() + result = genWM(params) + assert isinstance(result, tuple) + assert len(result) == 5 + assert result[0].shape == (6,) + assert result[1].shape == (97,) + assert result[2].shape == (41,) + assert result[3].shape == (155,) + assert result[4].shape == (155, 41, 97, 6) diff --git a/tests/rcm/wmutils/test_wmData.py b/tests/rcm/wmutils/test_wmData.py new file mode 100644 index 0000000..5da89ae --- /dev/null +++ b/tests/rcm/wmutils/test_wmData.py @@ -0,0 +1,30 @@ +import pytest +from kaipy.rcm.wmutils.wmData import wmParams + +def test_wmParams_initialization(): + params = wmParams(dim=5, nKp=7, nMLT=36, nL=40, nEk=100) + + assert params.dim == 5 + assert params.nKp == 7 + assert params.nMLT == 36 + assert params.nL == 40 + assert params.nEk == 100 + +def test_wmParams_default_initialization(): + params = wmParams() + + assert params.dim == 4 + assert params.nKp == 6 + assert params.nMLT == 97 + assert params.nL == 41 + assert params.nEk == 155 + +def test_wmParams_getAttrs(): + params = wmParams(dim=4, nKp=6, nMLT=97, nL=41, nEk=155) + attrs = params.getAttrs() + + assert attrs['tauDim'] == 4 + assert attrs['nKp'] == 6 + assert attrs['nMLT'] == 97 + assert attrs['nL'] == 41 + assert attrs['nEk'] == 155 \ No newline at end of file