Adds tests for RCM

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)

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

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()

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'

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
+

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]))
+
+
+
+

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)

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