Fitting models

.github/workflows/check-pr-main-version.yml

@@ -17,16 +17,24 @@ jobs:
 
     steps:
       - name: Checkout PR branch (head)
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
         with:
           ref: ${{ github.head_ref }}
           fetch-depth: 0
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v5
+        uses: astral-sh/setup-uv@v7
+
+      - name: Ensure uv.lock is up to date
+        run: |
+          if ! uv lock --check; then
+            echo "❌ uv.lock is out of date."
+            echo "Please update uv.lock before versioning can proceed."
+            exit 1
+          fi
 
       - name: Install the project
-        run: uv sync --all-extras --dev --locked
+        run: uv sync --locked --group test
 
       - name: Get PR branch version
         id: pr_version
@@ -86,7 +94,7 @@ jobs:
             echo "Version matches — bumping version and creating new release branch"
 
             uv version --bump patch
-            uv sync --locked
+            uv sync --locked --group test
             NEXT_VERSION=$(uv version --short)
 
             if [[ "$NEXT_VERSION" == "$PR_VERSION" ]]; then

.github/workflows/check-uv-lock.yml

@@ -0,0 +1,83 @@
+name: Check uv.lock is up to date
+
+on:
+  pull_request:
+    branches:
+      - main
+
+permissions:
+  contents: write
+  pull-requests: write
+
+jobs:
+  check-uv-lock:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout PR branch
+        uses: actions/checkout@v6
+        with:
+          ref: ${{ github.head_ref }}
+          fetch-depth: 0
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
+
+      - name: Check lock file
+        id: lockcheck
+        run: |
+          if uv lock --check; then
+            echo "up_to_date=true" >> $GITHUB_OUTPUT
+          else
+            echo "up_to_date=false" >> $GITHUB_OUTPUT
+          fi
+
+      - name: Exit early if lock is up to date
+        if: steps.lockcheck.outputs.up_to_date == 'true'
+        run: echo "uv.lock is up to date ✔"
+
+      - name: Regenerate lock file
+        if: steps.lockcheck.outputs.up_to_date == 'false'
+        run: |
+          if [[ "${GITHUB_HEAD_REF}" == release/* ]]; then
+            echo "❌ Release branch — uv.lock must already be up to date."
+            exit 1
+          fi
+
+          uv lock
+
+      - name: Commit updated lock file
+        if: steps.lockcheck.outputs.up_to_date == 'false'
+        run: |
+          SAFE_REF=$(echo "${GITHUB_HEAD_REF}" | tr '/' '-')
+
+          git config user.name "quantem-bot"
+          git config user.email "quantembot@gmail.com"
+
+          git checkout -b lockfix/${SAFE_REF}
+          git commit -am "Update uv.lock"
+          git push origin lockfix/${SAFE_REF}
+
+      - name: Create PR for updated lock file
+        if: steps.lockcheck.outputs.up_to_date == 'false'
+        env:
+          GH_TOKEN: ${{ secrets.QUANTEM_BOT_PAT }}
+        run: |
+          SAFE_REF=$(echo "${GITHUB_HEAD_REF}" | tr '/' '-')
+
+          if gh pr list --head "lockfix/${SAFE_REF}" --json number --jq 'length > 0'; then
+            echo "Lockfix PR already exists — skipping."
+            exit 1
+          fi
+
+          gh pr create \
+            --base "${{ github.base_ref }}" \
+            --head "lockfix/${SAFE_REF}" \
+            --title "Update uv.lock" \
+            --body "This PR updates \`uv.lock\` to match the current dependency specification."
+
+      - name: Fail original PR
+        if: steps.lockcheck.outputs.up_to_date == 'false'
+        run: |
+          echo "❌ uv.lock was out of date — a fix PR has been opened."
+          exit 1

.github/workflows/deploy.yml

@@ -13,10 +13,10 @@ jobs:
       UV_PUBLISH_TOKEN: ${{ secrets.PYPI_TOKEN }}
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v5
+        uses: astral-sh/setup-uv@v7
 
       - name: Verify lock file is up to date
         run: uv lock --check

.github/workflows/sync-dev-main.yml

@@ -14,7 +14,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout repo
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
 
       - name: Open PR from main to dev
         env:

.github/workflows/update-lock-file-cron.yml

@@ -13,12 +13,12 @@ jobs:
   update-deps:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
         with:
           ref: dev
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v5
+        uses: astral-sh/setup-uv@v7
 
       - name: Update dependencies
         env:

.github/workflows/uv-pytests.yml

@@ -13,23 +13,23 @@ jobs:
     strategy:
       matrix:
         # Test other OSs main branch PRs
-        os: ${{ github.base_ref == 'main' && fromJSON('["ubuntu-latest", "windows-latest", "macos-latest", "macos-15-intel"]') || fromJSON('["ubuntu-latest"]') }}
+        os: ${{ github.base_ref == 'main' && fromJSON('["ubuntu-latest", "windows-latest", "macos-latest"]') || fromJSON('["ubuntu-latest"]') }}
         python-version:
           - "3.11"
           - "3.12"
           - "3.13"
           - "3.14"
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
 
       - name: Install uv and set python version
-        uses: astral-sh/setup-uv@v5
+        uses: astral-sh/setup-uv@v7
         with:
           python-version: ${{ matrix.python-version }}
 
       - name: Install the project
-        run: uv sync --locked --all-extras --group test
+        run: uv sync --locked --group test
 
       - name: Run pytests
         run: uv run pytest tests

pyproject.toml

@@ -29,7 +29,7 @@ members = ["widget"]
 
 [project]
 name = "quantem"
-version = "0.1.7"
+version = "0.1.8"
 description = "quantitative electron microscopy analysis toolkit."
 keywords = ["EM","TEM","STEM","4DSTEM"]
 readme = "README.md"
@@ -53,6 +53,7 @@ dependencies = [
   "torchmetrics>=1.7.3",
   "optuna>=4.5.0",
   "hdf5plugin>=6.0.0",
+  "torchinfo>=1.8.0",
 ]
 
 [project.optional-dependencies]
@@ -72,6 +73,7 @@ only-include = [
 [dependency-groups]
 test = [
     "pytest>=8.3.5",
+    "pytest-cov>=7.0.0",
 ]
 dev = [
     { include-group = "test" },

src/quantem/core/config.py

@@ -270,6 +270,7 @@ def update(
 
     """
     for k, v in new.items():
+        k, v = check_key_val(k, v)
         k = canonical_name(k, old)
 
         if isinstance(v, Mapping):
@@ -480,11 +481,10 @@ def check_key_val(key: str, val: Any, deprecations: dict = deprecations) -> tupl
             new_val = "cpu"
         else:
             new_val, gpu_id = validate_device(new_val)
-            if config["has_torch"]:
-                if torch.cuda.is_available():
-                    torch.cuda.set_device(f"cuda:{gpu_id}")
-            if config["has_cupy"]:
-                cp.cuda.runtime.setDevice(gpu_id)
+            if "cuda" in new_val:
+                torch.cuda.set_device(gpu_id)
+                if config["has_cupy"]:
+                    cp.cuda.runtime.setDevice(gpu_id)
     return key, new_val
 
 
@@ -503,33 +503,42 @@ def validate_device(dev: str | int | torch.device | None = None) -> tuple[str, i
         dev = torch.device(
             "cuda" if torch.cuda.is_available() else "mps" if torch.mps.is_available() else "cpu"
         )
-    elif isinstance(dev, str) and dev.lower() == "gpu":
-        if torch.cuda.is_available():
-            dev = torch.device("cuda")
-        elif torch.mps.is_available():
-            dev = torch.device("mps")
-        else:
-            raise RuntimeError("Requested 'gpu' device, but neither CUDA nor MPS is available.")
-
-    # Convert to torch.device early
-    if isinstance(dev, (str, int)):
-        if isinstance(dev, int):
+    elif isinstance(dev, str):
+        if "cuda" in dev.lower():
+            dev = torch.device(dev)
+        elif "gpu" in dev.lower():
             if torch.cuda.is_available():
-                dev = torch.device(f"cuda:{dev}")
+                dev = torch.device("cuda")
             elif torch.mps.is_available():
                 dev = torch.device("mps")
             else:
-                dev = torch.device("cpu")
+                raise RuntimeError("gpu requested but cuda and mps are not available.")
+        elif dev.lower() == "mps":
+            dev = torch.device("mps")
+        elif dev.lower() == "cpu":
+            dev = torch.device("cpu")
         else:
-            dev = torch.device(dev)
+            raise ValueError(
+                f"Requested unknown device type: {dev} (must be 'cuda', 'mps', or 'cpu')"
+            )
+    elif isinstance(dev, int):
+        if dev < 0:
+            raise ValueError(f"Requested negative GPU index: {dev} (must be >= 0)")
+        if torch.cuda.is_available():
+            dev = torch.device(f"cuda:{dev}")
+        else:
+            raise RuntimeError(f"Requested GPU index '{dev}' device, but cuda is not available.")
     elif not isinstance(dev, torch.device):
         raise TypeError(f"Unsupported device type: {type(dev)} ({dev})")
 
-    # Normalize to supported device types
     if dev.type == "cuda":
         if not torch.cuda.is_available():
             raise RuntimeError("CUDA device requested but not available.")
         index = dev.index if dev.index is not None else torch.cuda.current_device()
+        if index >= NUM_DEVICES:
+            raise RuntimeError(
+                f"CUDA device index {index} is out of range for {NUM_DEVICES} available devices."
+            )
         return f"cuda:{index}", index
 
     elif dev.type == "mps":
@@ -560,7 +569,18 @@ def write(path: Path | str = PATH / "config.yaml") -> None:
         yaml.dump(config, f)
 
 
-def set_device(dev: str | int) -> None:
+def set_device(dev: str | int | "torch.device") -> None:
+    """Set the current device. Accepts a torch-style string, an integer index, or a
+    torch.device object.
+    Examples
+    --------
+    >>> set_device("cuda:0")
+    >>> set_device(0)
+    >>> set_device(torch.device("cuda:0"))
+    >>> set_device("mps")
+    >>> set_device("cpu")
+    >>> set_device("gpu")
+    """
     set({"device": dev})