initial claude setup

.claude/rules/uv.md

@@ -0,0 +1,11 @@
+---
+description: Enforce uv as the package manager for all Python operations
+globs:
+  - "**/*.py"
+  - "**/pyproject.toml"
+---
+
+- Always use `uv run` to execute commands — never bare `python`, `pytest`, `ruff`, or other tools.
+- When working in this repo (local development, CI, or scripts), do not use `pip install`; use `uv sync` (with `--extra` or `--all-groups` flags) to manage dependencies. End-user installation docs may still reference `pip`.
+- The `uv.lock` file is the source of truth for resolved dependency versions.
+- When adding dependencies, add them to `pyproject.toml` and run `uv sync`.

.claude/settings.json

@@ -0,0 +1,15 @@
+{
+  "hooks": {
+    "PostToolUse": [
+      {
+        "matcher": "Edit|Write",
+        "hooks": [
+          {
+            "type": "command",
+            "command": "FILE=$(jq -r '.tool_input.file_path'); if [[ \"$FILE\" == *.py ]]; then cd \"$CLAUDE_PROJECT_DIR\" && uv run ruff format \"$FILE\" --quiet 2>/dev/null; fi"
+          }
+        ]
+      }
+    ]
+  }
+}

.gitignore

@@ -44,7 +44,6 @@ share/python-wheels/
 .installed.cfg
 *.egg
 MANIFEST
-uv.lock
 
 # PyInstaller
 #  Usually these files are written by a python script from a template

.pre-commit-config.yaml

@@ -1,19 +1,25 @@
 repos:
-  # Ruff linting and formatting
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.14.4
+    rev: v0.14.13
     hooks:
       - id: ruff
-        args: [--fix]
+        args: [--fix, --exit-non-zero-on-fix]
       - id: ruff-format
 
-  # Pre-commit hooks
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v6.0.0
     hooks:
       - id: trailing-whitespace
       - id: end-of-file-fixer
       - id: check-yaml
-        args: [--unsafe] # this helps with some MkDocs YAML files
+        exclude: ^mkdocs\.yml$
       - id: check-added-large-files
+        args: ['--maxkb=99900']
       - id: check-merge-conflict
+
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.19.1
+    hooks:
+      - id: mypy
+        args: [--no-strict-optional, --ignore-missing-imports]
+        additional_dependencies: [types-PyYAML]

CLAUDE.md

@@ -0,0 +1,138 @@
+# LANfactory — Project Context for Claude
+
+## What is LANfactory?
+
+Lightweight Python package for training Likelihood Approximation Networks (LANs), Choice Probability Networks (CPNs), and Option Probability Networks (OPNs) using PyTorch or JAX/Flax backends. Trained networks are exported to ONNX format and uploaded to HuggingFace for consumption by HSSM. This package sits in the middle of the HSSM ecosystem: it depends on ssm-simulators for training data and produces the neural network artifacts that HSSM uses at inference time. For ecosystem-wide context, see the HSSMSpine repo.
+
+## Project Structure
+
+```
+src/lanfactory/                # Main package
+  cli/                         # Typer CLIs: jaxtrain, torchtrain, transform-onnx, upload-hf, download-hf
+  config/                      # Default network and training configs (LAN, CPN, OPN)
+  trainers/                    # Training implementations (torch_mlp.py, jax_mlp.py)
+  onnx/                        # PyTorch → ONNX export
+  hf/                          # HuggingFace Hub integration (upload, download, model cards)
+  utils/                       # Config save/load, MLflow utilities
+tests/                         # pytest suite (trainers, CLI, ONNX, HuggingFace, E2E)
+docs/                          # MkDocs documentation + tutorial notebooks
+notebooks/                     # Test notebooks
+```
+
+## Build & Tooling
+
+- **Build system:** setuptools (pure Python, no compiled extensions)
+- **Package manager:** uv (with `uv.lock`)
+- **Python:** >3.10, <3.14 (classifiers target 3.11, 3.12, 3.13)
+- **Linting:** ruff (line length 88, via pre-commit)
+- **Type checking:** mypy
+- **No system dependencies** — unlike ssm-simulators, this is pure Python + PyTorch/Flax
+
+## Common Commands
+
+```bash
+# Install all dependency groups (e.g. dev)
+uv sync --all-groups
+
+# Run tests
+uv run pytest tests/
+
+# Lint & format
+uv run ruff check src/lanfactory && uv run ruff format --check .
+
+# Build docs
+uv run mkdocs build
+uv run mkdocs serve
+
+# Train a network (PyTorch)
+uv run torchtrain --config-path <yaml> --training-data-folder <dir> --networks-path-base <dir>
+
+# Train a network (JAX)
+uv run jaxtrain --config-path <yaml> --training-data-folder <dir> --networks-path-base <dir>
+
+# Export PyTorch model to ONNX
+uv run transform-onnx --network-config-file config.pickle --state-dict-file model.pt \
+  --input-shape 6 --output-onnx-file model.onnx
+
+# Upload trained models to HuggingFace
+uv run upload-hf --model-folder <dir> --network-type lan --model-name ddm
+
+# Download models from HuggingFace
+uv run download-hf --network-type lan --model-name ddm --output-folder <dir>
+```
+
+## Key Architecture Patterns
+
+### Network Types
+
+| Type | Full Name | Output | Loss | Use Case |
+|------|-----------|--------|------|----------|
+| LAN | Likelihood Approximation Network | logprob | Huber | Log-likelihood approximation |
+| CPN | Choice Probability Network | logits | BCE with logits | Choice probability estimation |
+| OPN | Option Probability Network | logits | BCE with logits | Option probability estimation |
+
+All three use the same MLP architecture (`[100, 100, 1]` default, tanh activations)
+but differ in output type and loss function.
+
+### Training Backends
+
+- **PyTorch** (`torchtrain` CLI, `trainers/torch_mlp.py`) — primary backend.
+  Supports CUDA, ONNX export, full training loop with validation.
+- **JAX/Flax** (`jaxtrain` CLI, `trainers/jax_mlp.py`) — alternative backend.
+  Uses optax optimizers. No native ONNX export (train in JAX, convert via PyTorch if needed).
+
+### ONNX Export Pipeline
+
+PyTorch model → `torch.onnx.export()` → `.onnx` file. This is the format
+HSSM consumes at runtime. Only PyTorch models can be directly exported to ONNX.
+
+### HuggingFace Integration
+
+- **Upload:** `lanfactory.hf.upload_model()` — uploads `.onnx`, `.pt`, config pickles,
+  and auto-generated README to `franklab/HSSM` on HuggingFace.
+  Requires `model_card.yaml` in the model folder.
+- **Download:** `lanfactory.hf.download_model()` — downloads by network type + model name.
+- **Default repo:** `franklab/HSSM`
+- **Optional dependency:** `huggingface-hub>=0.20.0` (install via `uv sync --extra hf`)
+
+### Config System
+
+Training configs are YAML files parsed by the CLI. Key fields:
+- `NETWORK_TYPE`: `lan`, `cpn`, or `opn`
+- `layer_sizes`, `activations`: network architecture
+- `n_epochs`, `learning_rate`, `loss`, `optimizer`: training hyperparams
+- `cpu_batch_size`, `gpu_batch_size`: device-specific batch sizes
+
+Default configs available in `lanfactory.config.network_configs`.
+
+### MLflow Integration
+
+Optional experiment tracking via MLflow. CLI flags: `--mlflow-run-name`, `--mlflow-experiment-name`,
+`--mlflow-tracking-uri`, `--mlflow-artifact-location`. Supports resuming runs via `--mlflow-run-id`.
+
+## CLI Entry Points
+
+| Command | Module | Purpose |
+|---------|--------|---------|
+| `torchtrain` | `lanfactory.cli.torch_train` | Train PyTorch networks from YAML config |
+| `jaxtrain` | `lanfactory.cli.jax_train` | Train JAX networks from YAML config |
+| `transform-onnx` | `lanfactory.onnx.transform_onnx` | Convert PyTorch model → ONNX |
+| `upload-hf` | `lanfactory.cli.upload_hf` | Upload trained models to HuggingFace |
+| `download-hf` | `lanfactory.cli.download_hf` | Download models from HuggingFace |
+
+## CI Workflows
+
+| Workflow | Purpose |
+|----------|---------|
+| `run_tests.yml` | Tests on Python 3.11/3.12/3.13 + ruff lint/format + codecov |
+| `build_wheels.yml` | Build sdist, upload to TestPyPI → PyPI on release publish |
+
+## Known Issues
+
+- `__init__.py` version (`0.5.3`) is out of sync with `pyproject.toml` (`0.6.1`)
+
+## Compaction
+
+When compacting, preserve: file list of modified files, the three network types
+(LAN/CPN/OPN) and their differences, CLI entry points, ONNX export flow,
+HuggingFace upload/download interface, and all test commands.

pyproject.toml

@@ -72,6 +72,7 @@ dev = [
     "pytest-xdist>=3.6.1",
     "pytest>=8.3.1",
     "ruff>=0.14.4",
+    "types-PyYAML",
     "mlflow>=3.6.0",
 ]
 
@@ -94,11 +95,11 @@ addopts = [
 ]
 
 [tool.ruff]
-# Specify the directories to check
+line-length = 88
 src = ["src/lanfactory", "tests"]
 exclude = [
     "notebooks/*",
-    "docs/*",                                   # Exclude the entire docs folder
+    "docs/*",
     "docs/basic_tutorial/basic_tutorial.ipynb",
 ]
 
@@ -131,8 +132,6 @@ exclude_lines = [
     "logger.error",
 ]
 
-# Add any additional configuration options for ruff
-line-length = 120
 
 
 [project.scripts]

src/lanfactory/hf/upload.py

@@ -233,19 +233,17 @@ def _collect_files(
     list[Path]
         List of file paths to upload.
     """
-    files = set()
+    matched: set[Path] = set()
 
     # Collect files matching include patterns
     for pattern in include_patterns:
-        files.update(folder.glob(pattern))
+        matched.update(folder.glob(pattern))
 
     # Remove files matching exclude patterns
     if exclude_patterns:
         for pattern in exclude_patterns:
             excluded = set(folder.glob(pattern))
-            files -= excluded
+            matched -= excluded
 
     # Filter to only regular files (not directories)
-    files = [f for f in files if f.is_file()]
-
-    return sorted(files)
+    return sorted(f for f in matched if f.is_file())

src/lanfactory/trainers/torch_mlp.py

@@ -337,8 +337,9 @@ def TorchMLPFactory(
     """
     if isinstance(network_config, str):
         with open(network_config, "rb") as f:
-            network_config = pickle.load(f)
+            network_config = pickle.load(f)  # noqa: S301
 
+    assert isinstance(network_config, dict)
     return TorchMLP(
         network_config=network_config,
         input_shape=input_dim,
@@ -506,7 +507,7 @@ def __init__(
         elif isinstance(train_config, dict):
             print("train_config is passed as dictionary: \n")
             print(train_config)
-            self.train_config: dict = train_config
+            self.train_config = train_config
 
         self.model: TorchMLP = model.to(self.dev)
         self.allow_abs_path_folder_generation: bool = allow_abs_path_folder_generation