Merge pull request #57 from gridfm/improve_tests

Add data downloading functionality and gdown dependency for test setup

Author: romeokienzler

README.md

@@ -33,6 +33,8 @@ Install gridfm-graphkit in editable mode
 pip install -e .
 ```
 
+**`torch-scatter` is a required dependency.** It cannot be bundled in `pyproject.toml` because the correct wheel depends on your PyTorch and CUDA versions, so it must be installed separately.
+
 Get PyTorch + CUDA version for torch-scatter
 ```bash
 TORCH_CUDA_VERSION=$(python -c "import torch; print(torch.__version__ + ('+cpu' if torch.version.cuda is None else ''))")

docs/install/installation.md

@@ -15,6 +15,8 @@ Install gridfm-graphkit in editable mode
 pip install -e .
 ```
 
+**`torch-scatter` is a required dependency.** It cannot be bundled in `pyproject.toml` because the correct wheel depends on your PyTorch and CUDA versions, so it must be installed separately.
+
 Get PyTorch + CUDA version for torch-scatter
 
 ```bash

integrationtests/conftest.py

@@ -0,0 +1,31 @@
+import pytest
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--calibrate",
+        type=int,
+        default=0,
+        help="Run training N times to collect metric mean/std for range calibration. "
+        "Skips metric range assertions. Example: pytest --calibrate 5",
+    )
+    parser.addoption(
+        "--ci",
+        type=float,
+        default=0.995,
+        help="Confidence interval level for calibration stats (default 0.995). "
+        "Example: pytest --calibrate 5 -s --ci 0.995",
+    )
+
+
+@pytest.fixture
+def calibrate_runs(request):
+    """Number of calibration runs requested via --calibrate (0 = normal test mode)."""
+    return request.config.getoption("--calibrate")
+
+
+@pytest.fixture
+def ci_level(request):
+    """Confidence interval level requested via --ci (default 0.995)."""
+    return request.config.getoption("--ci")
+

integrationtests/generate_test_data.py

@@ -0,0 +1,72 @@
+import urllib.request
+import yaml
+import subprocess
+
+
+def execute_and_live_output(cmd) -> None:
+    subprocess.run(cmd, text=True, shell=True, check=True)
+
+
+def _base_config() -> dict:
+    """
+    Download the default config from gridfm-datakit and apply common test parameters.
+    """
+    config_url = (
+        "https://raw.githubusercontent.com/gridfm/gridfm-datakit/refs/heads/main"
+        "/scripts/config/default.yaml"
+    )
+
+    print(f"Downloading config from {config_url}...")
+    with urllib.request.urlopen(config_url) as response:
+        config_content = response.read().decode("utf-8")
+
+    config = yaml.safe_load(config_content)
+
+    config["network"]["name"] = "case14_ieee"
+    config["load"]["scenarios"] = 10000
+    config["topology_perturbation"]["n_topology_variants"] = 2
+
+    return config
+
+
+def generate_pf_test_data(config_path: str = "integrationtests/default_pf.yaml") -> None:
+    """
+    Generate power-flow (PF) test data for case14_ieee with 10 000 scenarios
+    and 2 topology variants.
+    """
+    config = _base_config()
+
+    with open(config_path, "w") as f:
+        yaml.dump(config, f, default_flow_style=False, sort_keys=False)
+
+    print(f"PF config written to {config_path}")
+    print(f"  network.name                          : {config['network']['name']}")
+    print(f"  load.scenarios                        : {config['load']['scenarios']}")
+    print(f"  topology_perturbation.n_topology_variants: {config['topology_perturbation']['n_topology_variants']}")
+
+    execute_and_live_output(f"gridfm_datakit generate {config_path}")
+
+
+def generate_opf_test_data(config_path: str = "integrationtests/default_opf.yaml") -> None:
+    """
+    Generate optimal power-flow (OPF) test data for case14_ieee with 10 000 scenarios
+    and 2 topology variants.
+    """
+    config = _base_config()
+    config["settings"]["mode"] = "opf"
+
+    with open(config_path, "w") as f:
+        yaml.dump(config, f, default_flow_style=False, sort_keys=False)
+
+    print(f"OPF config written to {config_path}")
+    print(f"  network.name                          : {config['network']['name']}")
+    print(f"  load.scenarios                        : {config['load']['scenarios']}")
+    print(f"  topology_perturbation.n_topology_variants: {config['topology_perturbation']['n_topology_variants']}")
+    print(f"  settings.mode                         : {config['settings']['mode']}")
+
+    execute_and_live_output(f"gridfm_datakit generate {config_path}")
+
+
+if __name__ == "__main__":
+    #generate_pf_test_data()
+    generate_opf_test_data()