Feature/refactor

.github/workflows/main.yml

@@ -11,7 +11,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        py-version: [ '3.8', '3.9', '3.10' ]
+        py-version: [ '3.9', '3.10', '3.11' ]
 
     steps:
       - uses: actions/checkout@v4
@@ -22,10 +22,10 @@ jobs:
           python-version: ${{ matrix.py-version }}
 
       - name: Install torchcast
-        run: pip install .[tests]
+        run: pip install .[tests] --extra-index-url https://download.pytorch.org/whl/cpu
 
       - name: Run tests
-        run: python3 -m unittest
+        run: pytest
 
   release:
     needs: test
@@ -74,7 +74,7 @@ jobs:
       - name: Install dependencies
         run: |
           pip install awscli
-          pip install ".[docs]"
+          pip install ".[docs]" --extra-index-url https://download.pytorch.org/whl/cpu
 
       - name: Install pandoc
         uses: pandoc/actions/setup@main

docs/quick_start.py

@@ -129,7 +129,7 @@
 
 # %%
 pred.plot(
-    pred.to_dataframe(dataset_all, type='components').query("group=='Changping'"), split_dt=SPLIT_DT,
+    pred.to_dataframe(dataset_all, type='observed_states').query("group=='Changping'"), split_dt=SPLIT_DT,
     time_colname='time', group_colname='group'
 )
 

pyproject.toml

@@ -2,14 +2,14 @@
 name = "torchcast"
 description = "Forecasting in PyTorch"
 readme = "README.md"
-requires-python = ">= 3.8"
+requires-python = ">= 3.9"
 classifiers = [
     "Development Status :: 5 - Production/Stable",
     "Intended Audience :: Developers",
     "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3.8",
     "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
 ]
 license = { file = "LICENSE" }
 authors = [
@@ -20,7 +20,8 @@ dependencies = [
     "numpy>=1.4",
     "scipy>=1.10",
     "tqdm>=4.59",
-    "platformdirs>=4.3"
+    "platformdirs>=4.3",
+    "pandas<3"
 ]
 dynamic = ["version"]
 
@@ -35,9 +36,8 @@ dev = [
     "black",
 ]
 tests = [
-    "parameterized>=0.7",
+    "pytest>=7.0",
     "filterpy>=1.4",
-    "pandas>=1.0"
 ]
 docs = [
     "jupytext>=1.11",

tests/test_covariance.py

@@ -1,27 +1,26 @@
 import torch
 from torchcast.covariance import Covariance
-import unittest
 
 
-class TestCovariance(unittest.TestCase):
-    torch.no_grad()
+@torch.no_grad()
+def test_from_log_cholesky():
+    module = Covariance(id='test', rank=3)
 
-    def test_from_log_cholesky(self):
-        module = Covariance(id='test', rank=3)
+    module.state_dict()['cholesky_log_diag'][:] = torch.arange(1., 3.1)
+    module.state_dict()['cholesky_off_diag'][:] = torch.arange(1., 3.1)
 
-        module.state_dict()['cholesky_log_diag'][:] = torch.arange(1., 3.1)
-        module.state_dict()['cholesky_off_diag'][:] = torch.arange(1., 3.1)
+    expected = torch.tensor([[7.3891, 2.7183, 5.4366],
+                             [2.7183, 55.5982, 24.1672],
+                             [5.4366, 24.1672, 416.4288]])
+    diff = (expected - module({}, num_groups=1, num_times=1)).abs()
+    assert (diff < .0001).all()
 
-        expected = torch.tensor([[7.3891, 2.7183, 5.4366],
-                                 [2.7183, 55.5982, 24.1672],
-                                 [5.4366, 24.1672, 416.4288]])
-        diff = (expected - module({}, num_groups=1, num_times=1)).abs()
-        self.assertTrue((diff < .0001).all())
 
-    def test_empty_idx(self):
-        module = torch.jit.script(Covariance(id='test', rank=3, empty_idx=[0]))
-        cov = module({}, num_groups=1, num_times=1)
-        cov = cov.squeeze()
-        self.assertTrue((cov[0, :] == 0).all())
-        self.assertTrue((cov[:, 0] == 0).all())
-        self.assertTrue((cov == cov.t()).all())
+@torch.no_grad()
+def test_empty_idx():
+    module = Covariance(id='test', rank=3, empty_idx=[0])
+    cov = module({}, num_groups=1, num_times=1)
+    cov = cov.squeeze()
+    assert (cov[0, :] == 0).all()
+    assert (cov[:, 0] == 0).all()
+    assert (cov == cov.t()).all()

tests/test_data_utils.py

@@ -1,4 +1,3 @@
-import unittest
 from warnings import warn
 
 import numpy as np
@@ -7,53 +6,53 @@
 from torchcast.utils.data import TimeSeriesDataset
 
 
-class TestDataUtils(unittest.TestCase):
-    def test_time_series_dataset(self):
-        values = torch.randn((3, 39, 2))
-
-        batch = TimeSeriesDataset(
-            values,
-            group_names=['one', 'two', 'three'],
-            start_times=[0, 0, 0],
-            measures=[['y1', 'y2']],
-            dt_unit=None
-        )
-        try:
-            import pandas as pd
-        except ImportError:
-            warn("Not testing TimeSeriesDataset.to_dataframe, pandas not installed.")
-            return
-        df1 = batch.to_dataframe()
-
-        df2 = pd.concat([
-            pd.DataFrame(values[i].numpy(), columns=batch.all_measures).assign(group=group, time=batch.times()[0])
-            for i, group in enumerate(batch.group_names)
-        ])
-        self.assertTrue((df1 == df2).all().all())
-
-    def test_pad_x(self, num_times: int = 10):
-        from pandas import DataFrame
-        df = DataFrame({'x1': np.random.randn(num_times), 'x2': np.random.randn(num_times)})
-        df['y'] = 1.5 * df['x1'] + -.5 * df['x2'] + .1 * np.random.randn(num_times)
-        df['time'] = df.index.values
-        df['group'] = '1'
-        dataset1 = TimeSeriesDataset.from_dataframe(
-            dataframe=df,
-            group_colname='group',
-            time_colname='time',
-            dt_unit=None,
-            X_colnames=['x1', 'x2'],
-            y_colnames=['y']
-        )
-        dataset2 = TimeSeriesDataset.from_dataframe(
-            dataframe=df,
-            group_colname='group',
-            time_colname='time',
-            dt_unit=None,
-            X_colnames=['x1', 'x2'],
-            y_colnames=['y'],
-            pad_X=None
-        )
-        self.assertFalse(torch.isnan(dataset1.tensors[1]).any())
-        self.assertFalse(torch.isnan(dataset2.tensors[1]).any())
-        self.assertTrue((dataset1.tensors[1] == dataset2.tensors[1]).all())
+def test_time_series_dataset():
+    values = torch.randn((3, 39, 2))
+
+    batch = TimeSeriesDataset(
+        values,
+        group_names=['one', 'two', 'three'],
+        start_times=[0, 0, 0],
+        measures=[['y1', 'y2']],
+        dt_unit=None
+    )
+    try:
+        import pandas as pd
+    except ImportError:
+        warn("Not testing TimeSeriesDataset.to_dataframe, pandas not installed.")
+        return
+    df1 = batch.to_dataframe()
+
+    df2 = pd.concat([
+        pd.DataFrame(values[i].numpy(), columns=batch.all_measures).assign(group=group, time=batch.times()[0])
+        for i, group in enumerate(batch.group_names)
+    ])
+    assert (df1 == df2).all().all()
+
+
+def test_pad_x(num_times: int = 10):
+    from pandas import DataFrame
+    df = DataFrame({'x1': np.random.randn(num_times), 'x2': np.random.randn(num_times)})
+    df['y'] = 1.5 * df['x1'] + -.5 * df['x2'] + .1 * np.random.randn(num_times)
+    df['time'] = df.index.values
+    df['group'] = '1'
+    dataset1 = TimeSeriesDataset.from_dataframe(
+        dataframe=df,
+        group_colname='group',
+        time_colname='time',
+        dt_unit=None,
+        X_colnames=['x1', 'x2'],
+        y_colnames=['y']
+    )
+    dataset2 = TimeSeriesDataset.from_dataframe(
+        dataframe=df,
+        group_colname='group',
+        time_colname='time',
+        dt_unit=None,
+        X_colnames=['x1', 'x2'],
+        y_colnames=['y'],
+        pad_X=None
+    )
+    assert not torch.isnan(dataset1.tensors[1]).any()
+    assert not torch.isnan(dataset2.tensors[1]).any()
+    assert (dataset1.tensors[1] == dataset2.tensors[1]).all()