Merge pull request #13 from SWxTREC/flattened-input-arrays

Handle flattened input arrays

Author: greglucas

.coveragerc

@@ -0,0 +1,4 @@
+[run]
+branch = true
+source = pymsis
+omit = pymsis/__init__.py

.github/workflows/ci.yml

@@ -53,8 +53,7 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install flake8 pytest
-        pip install -r requirements.txt
+        pip install -r requirements-test.txt
 
     - name: Install pymsis
       run: pip install -v .
@@ -65,12 +64,7 @@ jobs:
         flake8 . --count --show-source --statistics
 
     - name: Test with pytest
-      run: pytest --junitxml=junit/test-results-${{ matrix.python-version }}.xml
+      run: pytest --color=yes --cov --cov-report=xml
 
-    - name: Upload pytest test results
-      uses: actions/upload-artifact@v2
-      with:
-        name: pytest-results-${{ matrix.python-version }}
-        path: junit/test-results-${{ matrix.python-version }}.xml
-      # Use always() to always run this step to publish test results when there are test failures
-      if: ${{ always() }}
+    - name: Upload code coverage
+      uses: codecov/codecov-action@v2

codecov.yml

@@ -0,0 +1 @@
+comment: false

pymsis/msis.py

@@ -7,7 +7,13 @@
 
 def run(dates, lons, lats, alts, f107s, f107as, aps,
         options=None, version=2):
-    """Call MSIS looping over all possible inputs.
+    """
+    Call MSIS looping over all possible inputs. If ndates is
+    the same as nlons, nlats, and nalts, then a flattened
+    multi-point input array is assumed. Otherwise, the data
+    will be expanded in a grid-like fashion. The possible
+    return shapes are therefore (ndates, 11) and
+    (ndates, nlons, nlats, nalts, 11).
 
     Parameters
     ----------
@@ -32,7 +38,7 @@ def run(dates, lons, lats, alts, f107s, f107as, aps,
 
     Returns
     -------
-    ndarray (ndates, nlons, nlats, nalts, 11)
+    ndarray (ndates, nlons, nlats, nalts, 11) or (ndates, 11)
         | The data calculated at each grid point:
         | [Total mass density (kg/m3)
         | N2 # density (m-3),
@@ -168,7 +174,8 @@ def create_input(dates, lons, lats, alts, f107s, f107as, aps):
     (shape, flattened_input)
         The shape of the data as a tuple (ndates, nlons, nlats, nalts) and
         the flattened version of the input data
-        (ndates*nlons*nlats*nalts, 14).
+        (ndates*nlons*nlats*nalts, 14). If the input array was preflattened
+        (ndates == nlons == nlats == nalts), then the shape is (ndates,).
     """
     # Turn everything into arrays
     dates = np.atleast_1d(np.array(dates, dtype='datetime64'))
@@ -194,18 +201,30 @@ def create_input(dates, lons, lats, alts, f107s, f107as, aps):
     nlons = len(lons)
     nlats = len(lats)
     nalts = len(alts)
-    shape = (ndates, nlons, nlats, nalts)
 
     if not (ndates == len(f107s) == len(f107as) == len(aps)):
         raise ValueError(f"The length of dates ({ndates}), f107s "
                          f"({len(f107s)}), f107as ({len(f107as)}), "
                          f"and aps ({len(aps)}) must all be equal")
 
+    if ndates == nlons == nlats == nalts:
+        # This means the data came in preflattened, from a satellite
+        # trajectory for example, where we don't want to make a grid
+        # out of the input data, we just want to stack it together.
+        arr = np.stack([dyear, dseconds, lons, lats, alts, f107s, f107as], -1)
+
+        # ap has 7 components, so we need to concatenate it onto the
+        # arrays rather than stack
+        flattened_input = np.concatenate([arr, aps], axis=1,
+                                         dtype=np.float32)
+        shape = (ndates,)
+        return shape, flattened_input
+
     # Make a grid of indices
     indices = np.stack(np.meshgrid(np.arange(ndates),
-                                   np.arange(nlons),
-                                   np.arange(nlats),
-                                   np.arange(nalts), indexing='ij'),
+                       np.arange(nlons),
+                       np.arange(nlats),
+                       np.arange(nalts), indexing='ij'),
                        -1).reshape(-1, 4)
 
     # Now stack all of the arrays, indexing by the proper indices
@@ -215,5 +234,7 @@ def create_input(dates, lons, lats, alts, f107s, f107as, aps):
                     f107s[indices[:, 0]], f107as[indices[:, 0]]], -1)
     # ap has 7 components, so we need to concatenate it onto the
     # arrays rather than stack
-    return shape, np.concatenate([arr, aps[indices[:, 0], :]], axis=1,
-                                 dtype=np.float32)
+    flattened_input = np.concatenate([arr, aps[indices[:, 0], :]], axis=1,
+                                     dtype=np.float32)
+    shape = (ndates, nlons, nlats, nalts)
+    return shape, flattened_input

test-requirements.txt renamed to requirements-test.txt

@@ -1,3 +1,4 @@
 -r requirements.txt
 flake8
 pytest
+pytest-cov

tests/test_msis.py

@@ -60,7 +60,7 @@ def test_create_options():
 
 def test_create_input_single_point(input_data, expected_input):
     shape, data = msis.create_input(*input_data)
-    assert shape == (1, 1, 1, 1)
+    assert shape == (1,)
     assert data.shape == (1, 14)
     assert_array_equal(data[0, :], expected_input)
 
@@ -70,7 +70,7 @@ def test_create_input_datetime(input_data, expected_input):
     # .item() gets the datetime object from the np.datetime64 object
     input_data = (input_data[0].item(),) + input_data[1:]
     shape, data = msis.create_input(*input_data)
-    assert shape == (1, 1, 1, 1)
+    assert shape == (1,)
     assert data.shape == (1, 14)
     assert_array_equal(data[0, :], expected_input)
 
@@ -144,11 +144,11 @@ def test_run_options(input_data):
 
 def test_run_single_point(input_data, expected_output):
     output = msis.run(*input_data)
-    assert output.shape == (1, 1, 1, 1, 11)
+    assert output.shape == (1, 11)
     assert_allclose(np.squeeze(output), expected_output, rtol=1e-5)
 
 
-def test_run_multi_point(input_data, expected_output):
+def test_run_gridded_multi_point(input_data, expected_output):
     date, lon, lat, alt, f107, f107a, ap = input_data
     # 5 x 5 surface
     input_data = (date, [lon]*5, [lat]*5, alt, f107, f107a, ap)
@@ -158,6 +158,19 @@ def test_run_multi_point(input_data, expected_output):
     assert_allclose(np.squeeze(output), expected, rtol=1e-5)
 
 
+def test_run_multi_point(input_data, expected_output):
+    # test multi-point run, like a satellite fly-through
+    # and make sure we don't grid the input data
+    # 5 input points
+    date, lon, lat, alt, f107, f107a, ap = input_data
+    input_data = ([date]*5, [lon]*5, [lat]*5, [alt]*5,
+                  [f107]*5, [f107a]*5, ap*5)
+    output = msis.run(*input_data)
+    assert output.shape == (5, 11)
+    expected = np.tile(expected_output, (5, 1))
+    assert_allclose(np.squeeze(output), expected, rtol=1e-5)
+
+
 def test_run_wrapped_lon(input_data, expected_output):
     date, _, lat, alt, f107, f107a, ap = input_data
 
@@ -204,7 +217,7 @@ def test_run_versions(input_data):
 
 def test_run_version00(input_data, expected_output00):
     output = msis.run(*input_data, version=0)
-    assert output.shape == (1, 1, 1, 1, 11)
+    assert output.shape == (1, 11)
     assert_allclose(np.squeeze(output), expected_output00, rtol=1e-5)