Use uv, ruff, mypy. Add workflows

.github/workflows/lint.yml

@@ -0,0 +1,29 @@
+name: lint
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  check:
+    runs-on: ubuntu-24.04
+    steps:
+      - uses: actions/checkout@v5
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version-file: ".python-version"
+      - name: Install uv
+        uses: astral-sh/setup-uv@v6
+      - name: Install project
+        run: uv sync --locked --dev
+      - name: Run ruff
+        run: uv run ruff check
+      - name: Run black
+        run: uv run black --check .
+      - name: Run mypy
+        run: uv run make mypy

.github/workflows/tests.yml

@@ -0,0 +1,25 @@
+name: tests
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  check:
+    runs-on: ubuntu-24.04
+    steps:
+      - uses: actions/checkout@v5
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version-file: ".python-version"
+      - name: Install uv
+        uses: astral-sh/setup-uv@v6
+      - name: Install project
+        run: uv sync --locked --dev --extra tensorflow
+      - name: Run tests
+        run: uv run make -C tests all_tf -j

.python-version

@@ -0,0 +1 @@
+3.12

Makefile

@@ -1,3 +1,10 @@
+RUFF_FLAGS =
+BLACK_FLAGS =
+MYPY_FLAGS =
+
+MYPY_TARGETS = \
+    src/odil/history.py \
+
 default:
 
 data:
@@ -6,4 +13,11 @@ data:
 release:
 	V=$$(sed -rn 's/^version = "(.*)"$$/\1/p' pyproject.toml) && git archive --prefix="odil-$$V/" -o "odil-$$V.tar.gz" HEAD
 
-.PHONY: default release
+lint:
+	ruff check --fix $(RUFF_FLAGS) .
+	black $(BLACK_FLAGS) .
+
+mypy:
+	mypy $(MYPY_FLAGS) $(MYPY_TARGETS)
+
+.PHONY: default release lint mypy

README.md

@@ -29,6 +29,14 @@ or
 pip install git+https://github.com/cselab/odil.git
 ```
 
+### Using `uv`
+
+```
+uv venv --python 3.12
+. .venv/bin/activate
+uv sync --group dev --extra tensorflow --extra jax
+```
+
 ## Using GPU
 
 To enable GPU support, provide a non-empty list of devices in `CUDA_VISIBLE_DEVICES`.

examples/basic/fields.ipynb

@@ -18,138 +18,131 @@
    "outputs": [],
    "source": [
     "import sys\n",
-    "if 'google.colab' in sys.modules and not 'odil' in sys.modules:\n",
+    "\n",
+    "if \"google.colab\" in sys.modules and \"odil\" not in sys.modules:\n",
     "    %pip install -q git+https://github.com/cselab/odil\n",
     "import argparse\n",
+    "\n",
     "import numpy as np\n",
+    "\n",
     "import odil\n",
-    "from odil import printlog\n",
+    "\n",
+    "\n",
     "def operator(ctx):\n",
     "    res = []\n",
     "\n",
     "    def func(x, y):\n",
     "        return x * 0.25 + y * 0.5\n",
     "\n",
     "    # Cells.\n",
-    "    xc, yc = ctx.points(loc='cc')\n",
-    "    uc = ctx.field('uc')\n",
-    "    res += [('uc', uc - func(xc, yc))]\n",
+    "    xc, yc = ctx.points(loc=\"cc\")\n",
+    "    uc = ctx.field(\"uc\")\n",
+    "    res += [(\"uc\", uc - func(xc, yc))]\n",
     "\n",
     "    # Nodes.\n",
-    "    xn, yn = ctx.points(loc='nn')\n",
-    "    un = ctx.field('un')\n",
-    "    res += [('un', un - func(xn, yn))]\n",
+    "    xn, yn = ctx.points(loc=\"nn\")\n",
+    "    un = ctx.field(\"un\")\n",
+    "    res += [(\"un\", un - func(xn, yn))]\n",
     "\n",
     "    # Faces in x.\n",
-    "    xfx, yfx = ctx.points(loc='nc')\n",
-    "    ufx = ctx.field('ufx')\n",
-    "    res += [('ufx', ufx - func(xfx, yfx))]\n",
+    "    xfx, yfx = ctx.points(loc=\"nc\")\n",
+    "    ufx = ctx.field(\"ufx\")\n",
+    "    res += [(\"ufx\", ufx - func(xfx, yfx))]\n",
     "\n",
     "    # Faces in y.\n",
-    "    xfy, yfy = ctx.points(loc='cn')\n",
-    "    ufy = ctx.field('ufy')\n",
-    "    res += [('ufy', ufy - func(xfy, yfy))]\n",
+    "    xfy, yfy = ctx.points(loc=\"cn\")\n",
+    "    ufy = ctx.field(\"ufy\")\n",
+    "    res += [(\"ufy\", ufy - func(xfy, yfy))]\n",
     "    return res\n",
     "\n",
     "\n",
     "def parse_args():\n",
-    "    parser = argparse.ArgumentParser(\n",
-    "        formatter_class=argparse.ArgumentDefaultsHelpFormatter)\n",
-    "    parser.add_argument('--Nx', type=int, default=8, help=\"Grid size in x\")\n",
-    "    parser.add_argument('--Ny', type=int, default=4, help=\"Grid size in y\")\n",
-    "    parser.add_argument('--plot', type=int, default=1, help=\"Plot fields\")\n",
+    "    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)\n",
+    "    parser.add_argument(\"--Nx\", type=int, default=8, help=\"Grid size in x\")\n",
+    "    parser.add_argument(\"--Ny\", type=int, default=4, help=\"Grid size in y\")\n",
+    "    parser.add_argument(\"--plot\", type=int, default=1, help=\"Plot fields\")\n",
     "    odil.util.add_arguments(parser)\n",
     "    odil.linsolver.add_arguments(parser)\n",
-    "    parser.set_defaults(outdir='out_fields')\n",
+    "    parser.set_defaults(outdir=\"out_fields\")\n",
     "    parser.set_defaults(multigrid=1)\n",
-    "    parser.set_defaults(optimizer='adamn', lr=1e-2)\n",
-    "    parser.set_defaults(frames=1,\n",
-    "                        plot_every=100,\n",
-    "                        report_every=10,\n",
-    "                        history_every=10)\n",
+    "    parser.set_defaults(optimizer=\"adamn\", lr=1e-2)\n",
+    "    parser.set_defaults(frames=1, plot_every=100, report_every=10, history_every=10)\n",
     "    return parser.parse_args()\n",
     "\n",
     "\n",
     "def plot(problem, state, epoch, frame, cbinfo=None):\n",
-    "    from odil import plotutil\n",
     "    import matplotlib.pyplot as plt\n",
+    "\n",
     "    domain = problem.domain\n",
     "    fig, ax = plt.subplots(figsize=(4, 2))\n",
-    "    kw = dict(vmin=0, vmax=1, cmap='Greys', clip_on=False, lw=0.5)\n",
+    "    kw = dict(vmin=0, vmax=1, cmap=\"Greys\", clip_on=False, lw=0.5)\n",
     "\n",
     "    # Cells.\n",
-    "    xc, yc = domain.points(loc='cc')\n",
-    "    uc = domain.field(state, 'uc')\n",
-    "    ax.scatter(xc, yc, s=10, c=uc, edgecolor='C0', label='uc', **kw)\n",
+    "    xc, yc = domain.points(loc=\"cc\")\n",
+    "    uc = domain.field(state, \"uc\")\n",
+    "    ax.scatter(xc, yc, s=10, c=uc, edgecolor=\"C0\", label=\"uc\", **kw)\n",
     "\n",
     "    # Nodes.\n",
-    "    xn, yn = domain.points(loc='nn')\n",
-    "    un = domain.field(state, 'un')\n",
-    "    ax.scatter(xn, yn, s=10, c=un, edgecolor='C1', label='un', **kw)\n",
+    "    xn, yn = domain.points(loc=\"nn\")\n",
+    "    un = domain.field(state, \"un\")\n",
+    "    ax.scatter(xn, yn, s=10, c=un, edgecolor=\"C1\", label=\"un\", **kw)\n",
     "\n",
     "    # Faces in x.\n",
-    "    xfx, yfx = domain.points(loc='nc')\n",
-    "    ufx = domain.field(state, 'ufx')\n",
-    "    ax.scatter(xfx, yfx, s=10, c=ufx, edgecolor='C2', label='ufx', **kw)\n",
+    "    xfx, yfx = domain.points(loc=\"nc\")\n",
+    "    ufx = domain.field(state, \"ufx\")\n",
+    "    ax.scatter(xfx, yfx, s=10, c=ufx, edgecolor=\"C2\", label=\"ufx\", **kw)\n",
     "\n",
     "    # Faces in y.\n",
-    "    xfy, yfy = domain.points(loc='cn')\n",
-    "    ufy = domain.field(state, 'ufy')\n",
-    "    ax.scatter(xfy, yfy, s=10, c=ufy, edgecolor='C3', label='ufy', **kw)\n",
-    "\n",
-    "    ax.legend(loc='lower left',\n",
-    "              bbox_to_anchor=(0.1, 1),\n",
-    "              ncol=4,\n",
-    "              handletextpad=0)\n",
-    "\n",
-    "    ax.pcolormesh(xn,\n",
-    "                  yn,\n",
-    "                  uc,\n",
-    "                  edgecolor='k',\n",
-    "                  shading='flat',\n",
-    "                  zorder=0,\n",
-    "                  **dict(kw, lw=0.5))\n",
-    "\n",
-    "    ax.set_aspect('equal')\n",
+    "    xfy, yfy = domain.points(loc=\"cn\")\n",
+    "    ufy = domain.field(state, \"ufy\")\n",
+    "    ax.scatter(xfy, yfy, s=10, c=ufy, edgecolor=\"C3\", label=\"ufy\", **kw)\n",
+    "\n",
+    "    ax.legend(loc=\"lower left\", bbox_to_anchor=(0.1, 1), ncol=4, handletextpad=0)\n",
+    "\n",
+    "    ax.pcolormesh(xn, yn, uc, edgecolor=\"k\", shading=\"flat\", zorder=0, **dict(kw, lw=0.5))\n",
+    "\n",
+    "    ax.set_aspect(\"equal\")\n",
     "    ax.set_axis_off()\n",
     "    display(fig)\n",
     "\n",
     "\n",
     "def make_problem(args):\n",
     "    dtype = np.float64 if args.double else np.float32\n",
-    "    domain = odil.Domain(cshape=(args.Nx, args.Ny),\n",
-    "                         dimnames=['x', 'y'],\n",
-    "                         lower=(0, 0),\n",
-    "                         upper=(2, 1),\n",
-    "                         dtype=dtype,\n",
-    "                         multigrid=args.multigrid,\n",
-    "                         mg_interp=args.mg_interp,\n",
-    "                         mg_axes=[True, True],\n",
-    "                         mg_nlvl=args.nlvl)\n",
-    "    xx, yy = domain.points('x', 'y', loc='cn')\n",
-    "    ixx, iyy = domain.indices('x', 'y', loc='cc')\n",
+    "    domain = odil.Domain(\n",
+    "        cshape=(args.Nx, args.Ny),\n",
+    "        dimnames=[\"x\", \"y\"],\n",
+    "        lower=(0, 0),\n",
+    "        upper=(2, 1),\n",
+    "        dtype=dtype,\n",
+    "        multigrid=args.multigrid,\n",
+    "        mg_interp=args.mg_interp,\n",
+    "        mg_axes=[True, True],\n",
+    "        mg_nlvl=args.nlvl,\n",
+    "    )\n",
+    "    xx, yy = domain.points(\"x\", \"y\", loc=\"cn\")\n",
+    "    ixx, iyy = domain.indices(\"x\", \"y\", loc=\"cc\")\n",
     "\n",
     "    from odil import Field\n",
     "\n",
     "    state = odil.State(\n",
     "        fields={\n",
-    "            'uc': Field(np.zeros(domain.size(loc='cc')), loc='cc'),\n",
-    "            'un': Field(np.zeros(domain.size(loc='nn')), loc='nn'),\n",
-    "            'ufx': Field(np.zeros(domain.size(loc='nc')), loc='nc'),\n",
-    "            'ufy': Field(np.zeros(domain.size(loc='cn')), loc='cn'),\n",
-    "            'net': domain.make_neural_net([2, 4, 2]),\n",
-    "        })\n",
+    "            \"uc\": Field(np.zeros(domain.size(loc=\"cc\")), loc=\"cc\"),\n",
+    "            \"un\": Field(np.zeros(domain.size(loc=\"nn\")), loc=\"nn\"),\n",
+    "            \"ufx\": Field(np.zeros(domain.size(loc=\"nc\")), loc=\"nc\"),\n",
+    "            \"ufy\": Field(np.zeros(domain.size(loc=\"cn\")), loc=\"cn\"),\n",
+    "            \"net\": domain.make_neural_net([2, 4, 2]),\n",
+    "        }\n",
+    "    )\n",
     "    state = domain.init_state(state)\n",
     "    problem = odil.Problem(operator, domain)\n",
     "    return problem, state\n",
     "\n",
-    "sys.argv = ['fields']\n",
+    "\n",
+    "sys.argv = [\"fields\"]\n",
     "args = parse_args()\n",
     "odil.setup_outdir(args)\n",
     "problem, state = make_problem(args)\n",
-    "callback = odil.make_callback(problem,\n",
-    "                                  args,\n",
-    "                                  plot_func=plot if args.plot else None)\n",
+    "callback = odil.make_callback(problem, args, plot_func=plot if args.plot else None)\n",
     "odil.util.optimize_grad(args, args.optimizer, problem, state, callback)"
    ]
   }

examples/basic/fields.py

@@ -1,11 +1,12 @@
 #!/usr/bin/env python3
 
 import argparse
+
+import matplotlib.pyplot as plt
 import numpy as np
+
 import odil
-from odil import printlog
-from odil import plotutil
-import matplotlib.pyplot as plt
+from odil import plotutil, printlog
 
 """
 Demonstrates the use of fields with values centered in cell, faces, and nodes.
@@ -56,8 +57,6 @@ def parse_args():
 
 
 def plot(problem, state, epoch, frame, cbinfo=None):
-    from odil import plotutil
-    import matplotlib.pyplot as plt
 
     domain = problem.domain
     fig, ax = plt.subplots(figsize=(4, 2))

examples/heat/README.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+# ruff: noqa W291
+
 import os
 
 text = """<!--- DO NOT EDIT THIS FILE. Generated by [[GEN]]-->
@@ -36,7 +38,7 @@ def add_subdir(subdir, title, cmd, img=True):
 add_subdir("out_odil", "ODIL Adam", "case=2 outdir=out_odil ./run")
 add_subdir("out_pinn", "PINN Adam", "case=2p gpus=0 outdir=out_pinn ./run")
 
-text += f"""### Training history
+text += """### Training history
 
 ```
 ./plot_train.py

examples/heat/heat.py

@@ -1,13 +1,14 @@
 #!/usr/bin/env python3
 
 import argparse
-import numpy as np
+import os
 import pickle
 
-import odil
-from odil import plotutil
 import matplotlib.pyplot as plt
-from odil import printlog
+import numpy as np
+
+import odil
+from odil import history, printlog
 from odil.runtime import tf
 
 
@@ -106,7 +107,7 @@ def apply_bc_u(st):
     res = [("fu", fu)]
 
     if extra.imp_size:
-        u = u_st[0]
+        u_st[0]
         # Rescale weight to the total number of points.
         k = args.kimp * (np.prod(ctx.size()) / extra.imp_size) ** 0.5
         fuimp = extra.imp_mask * (u_st[0][0] - extra.imp_u) * k
@@ -515,7 +516,7 @@ def add_extra(d, *keys):
         state.fields["u_net"] = domain.make_neural_net([2] + args.arch_u + [1])
         operator = operator_pinn
     else:
-        raise RuntimeError(f"Unknown solver={solver}")
+        raise RuntimeError(f"Unknown solver={args.solver}")
 
     if args.infer_k:
         state.fields["k_net"] = domain.make_neural_net([1] + args.arch_k + [1])
@@ -550,7 +551,7 @@ def main():
     )
     odil.util.optimize(args, args.optimizer, problem, state, callback)
 
-    with open("done", "w") as f:
+    with open("done", "w"):
         pass