diff --git a/.gitea/workflows/ci.yaml b/.gitea/workflows/ci.yaml new file mode 100644 index 0000000..2dc42d0 --- /dev/null +++ b/.gitea/workflows/ci.yaml @@ -0,0 +1,65 @@ +name: CI + +on: + push: + branches: [main] + pull_request: + branches: [main] + +jobs: + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install dependencies + run: | + pip install ruff mypy + pip install -e ".[dev]" || pip install ruff mypy numpy + + - name: Ruff check + run: ruff check solver/ freecad/ tests/ scripts/ + + - name: Ruff format check + run: ruff format --check solver/ freecad/ tests/ scripts/ + + type-check: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install dependencies + run: | + pip install mypy numpy + pip install torch --index-url https://download.pytorch.org/whl/cpu + pip install torch-geometric + pip install -e ".[dev]" + + - name: Mypy + run: mypy solver/ freecad/ + + test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install dependencies + run: | + pip install torch --index-url https://download.pytorch.org/whl/cpu + pip install torch-geometric + pip install -e ".[train,dev]" + + - name: Run tests + run: pytest tests/ freecad/tests/ -v --tb=short diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..67e9723 --- /dev/null +++ b/.gitignore @@ -0,0 +1,48 @@ +# Python +__pycache__/ +*.py[cod] +*$py.class +*.egg-info/ +dist/ +build/ +*.egg + +# Virtual environments +.venv/ +venv/ + +# IDE +.vscode/ +.idea/ +*.swp +*.swo +*~ + +# mypy / ruff / pytest +.mypy_cache/ +.ruff_cache/ +.pytest_cache/ + +# Data (large files tracked separately) +data/synthetic/*.pt +data/fusion360/*.json +data/fusion360/*.step +data/processed/*.pt +!data/**/.gitkeep + +# Model checkpoints +*.ckpt +*.pth +*.onnx +*.torchscript + +# Experiment tracking +wandb/ +runs/ + +# OS +.DS_Store +Thumbs.db + +# Environment +.env diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..69b8611 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,23 @@ +repos: + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.3.4 + hooks: + - id: ruff + args: [--fix] + - id: ruff-format + + - repo: https://github.com/pre-commit/mirrors-mypy + rev: v1.8.0 + hooks: + - id: mypy + additional_dependencies: + - torch>=2.2 + - numpy>=1.26 + args: [--ignore-missing-imports] + + - repo: https://github.com/compilerla/conventional-pre-commit + rev: v3.1.0 + hooks: + - id: conventional-pre-commit + stages: [commit-msg] + args: [feat, fix, docs, style, refactor, perf, test, build, ci, chore, revert] diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..ebbca0a --- /dev/null +++ b/Dockerfile @@ -0,0 +1,61 @@ +FROM nvidia/cuda:12.4.1-devel-ubuntu22.04 AS base + +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 + +# System deps +RUN apt-get update && apt-get install -y --no-install-recommends \ + python3.11 python3.11-venv python3.11-dev python3-pip \ + git wget curl \ + # FreeCAD headless deps + freecad \ + libgl1-mesa-glx libglib2.0-0 \ + && rm -rf /var/lib/apt/lists/* + +RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.11 1 + +# Create venv +RUN python -m venv /opt/venv +ENV PATH="/opt/venv/bin:$PATH" + +# Install PyTorch with CUDA +RUN pip install --no-cache-dir \ + torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124 + +# Install PyG +RUN pip install --no-cache-dir \ + torch-geometric \ + pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv \ + -f https://data.pyg.org/whl/torch-2.4.0+cu124.html + +WORKDIR /workspace + +# Install project +COPY pyproject.toml . +RUN pip install --no-cache-dir -e ".[train,dev]" || true + +COPY . . +RUN pip install --no-cache-dir -e ".[train,dev]" + +# ------------------------------------------------------------------- +FROM base AS cpu + +# CPU-only variant (for CI and non-GPU environments) +FROM python:3.11-slim AS cpu-only + +ENV PYTHONUNBUFFERED=1 + +RUN apt-get update && apt-get install -y --no-install-recommends \ + git freecad libgl1-mesa-glx libglib2.0-0 \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /workspace + +COPY pyproject.toml . +RUN pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu +RUN pip install --no-cache-dir torch-geometric + +COPY . . +RUN pip install --no-cache-dir -e ".[train,dev]" + +CMD ["pytest", "tests/", "-v"] diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..18d1309 --- /dev/null +++ b/Makefile @@ -0,0 +1,48 @@ +.PHONY: train test lint data-gen export format type-check install dev clean help + +PYTHON ?= python +PYTEST ?= pytest +RUFF ?= ruff +MYPY ?= mypy + +help: ## Show this help + @grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | sort | \ + awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-20s\033[0m %s\n", $$1, $$2}' + +install: ## Install core dependencies + pip install -e . + +dev: ## Install all dependencies including dev tools + pip install -e ".[train,dev]" + pre-commit install + pre-commit install --hook-type commit-msg + +train: ## Run training (pass CONFIG=path/to/config.yaml) + $(PYTHON) -m solver.training.train $(if $(CONFIG),--config-path $(CONFIG)) + +test: ## Run test suite + $(PYTEST) tests/ freecad/tests/ -v --tb=short + +lint: ## Run ruff linter + $(RUFF) check solver/ freecad/ tests/ scripts/ + +format: ## Format code with ruff + $(RUFF) format solver/ freecad/ tests/ scripts/ + $(RUFF) check --fix solver/ freecad/ tests/ scripts/ + +type-check: ## Run mypy type checker + $(MYPY) solver/ freecad/ + +data-gen: ## Generate synthetic dataset (pass CONFIG=path/to/config.yaml) + $(PYTHON) scripts/generate_synthetic.py $(if $(CONFIG),--config-path $(CONFIG)) + +export: ## Export trained model for deployment + $(PYTHON) export/package_model.py $(if $(MODEL),--model $(MODEL)) + +clean: ## Remove build artifacts and caches + rm -rf build/ dist/ *.egg-info/ + rm -rf .mypy_cache/ .pytest_cache/ .ruff_cache/ + find . -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null || true + find . -type f -name "*.pyc" -delete 2>/dev/null || true + +check: lint type-check test ## Run all checks (lint, type-check, test) diff --git a/README.md b/README.md index e69de29..761eb0b 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,71 @@ +# kindred-solver + +Assembly constraint prediction via GNN. Produces a trained model embedded in a FreeCAD workbench (Kindred Create library), later integrated into vanilla Create. + +## Overview + +`kindred-solver` predicts whether assembly constraints (joints) are independent or redundant using graph neural networks. Given an assembly graph where bodies are nodes and joints are edges, the model classifies each constraint and reports degrees of freedom per body. + +## Repository Structure + +``` +kindred-solver/ +├── solver/ # Core library +│ ├── datagen/ # Synthetic data generation (pebble game) +│ ├── datasets/ # PyG dataset adapters +│ ├── models/ # GNN architectures (GIN, GAT, NNConv) +│ ├── training/ # Training loops and configs +│ ├── evaluation/ # Metrics and visualization +│ └── inference/ # Runtime prediction API +├── freecad/ # FreeCAD integration +│ ├── workbench/ # FreeCAD workbench addon +│ ├── bridge/ # FreeCAD <-> solver interface +│ └── tests/ # Integration tests +├── export/ # Model packaging for Create +├── configs/ # Hydra configs (dataset, model, training, export) +├── scripts/ # CLI utilities +├── data/ # Datasets (not committed) +├── tests/ # Unit and integration tests +└── docs/ # Documentation +``` + +## Setup + +### Install (development) + +```bash +pip install -e ".[train,dev]" +pre-commit install +pre-commit install --hook-type commit-msg +``` + +### Using Make + +```bash +make help # show all targets +make dev # install all deps + pre-commit hooks +make test # run tests +make lint # run ruff linter +make type-check # run mypy +make check # lint + type-check + test +make train # run training +make data-gen # generate synthetic data +make export # export model +``` + +### Using Docker + +```bash +# GPU training +docker compose up train + +# Run tests (CPU) +docker compose up test + +# Generate data +docker compose up data-gen +``` + +## License + +Apache 2.0 diff --git a/configs/dataset/fusion360.yaml b/configs/dataset/fusion360.yaml new file mode 100644 index 0000000..68b7b21 --- /dev/null +++ b/configs/dataset/fusion360.yaml @@ -0,0 +1,12 @@ +# Fusion 360 Gallery dataset config +name: fusion360 +data_dir: data/fusion360 +output_dir: data/processed + +splits: + train: 0.8 + val: 0.1 + test: 0.1 + +stratify_by: complexity +seed: 42 diff --git a/configs/dataset/synthetic.yaml b/configs/dataset/synthetic.yaml new file mode 100644 index 0000000..8450ad8 --- /dev/null +++ b/configs/dataset/synthetic.yaml @@ -0,0 +1,24 @@ +# Synthetic dataset generation config +name: synthetic +num_assemblies: 100000 +output_dir: data/synthetic + +complexity_distribution: + simple: 0.4 # 2-5 bodies + medium: 0.4 # 6-15 bodies + complex: 0.2 # 16-50 bodies + +body_count: + min: 2 + max: 50 + +templates: + - chain + - tree + - loop + - star + - mixed + +grounded_ratio: 0.5 +seed: 42 +num_workers: 4 diff --git a/configs/export/production.yaml b/configs/export/production.yaml new file mode 100644 index 0000000..1be79ec --- /dev/null +++ b/configs/export/production.yaml @@ -0,0 +1,25 @@ +# Production model export config +model_checkpoint: checkpoints/finetune/best_val_loss.ckpt +output_dir: export/ + +formats: + onnx: + enabled: true + opset_version: 17 + dynamic_axes: true + torchscript: + enabled: true + +model_card: + version: "0.1.0" + architecture: baseline + training_data: + - synthetic_100k + - fusion360_gallery + +size_budget_mb: 50 + +inference: + device: cpu + batch_size: 1 + confidence_threshold: 0.8 diff --git a/configs/model/baseline.yaml b/configs/model/baseline.yaml new file mode 100644 index 0000000..ebefe9b --- /dev/null +++ b/configs/model/baseline.yaml @@ -0,0 +1,24 @@ +# Baseline GIN model config +name: baseline +architecture: gin + +encoder: + num_layers: 3 + hidden_dim: 128 + dropout: 0.1 + +node_features_dim: 22 +edge_features_dim: 22 + +heads: + edge_classification: + enabled: true + hidden_dim: 64 + graph_classification: + enabled: true + num_classes: 4 # rigid, under, over, mixed + joint_type: + enabled: true + num_classes: 12 + dof_regression: + enabled: true diff --git a/configs/model/gat.yaml b/configs/model/gat.yaml new file mode 100644 index 0000000..6592bc1 --- /dev/null +++ b/configs/model/gat.yaml @@ -0,0 +1,28 @@ +# Advanced GAT model config +name: gat_solver +architecture: gat + +encoder: + num_layers: 4 + hidden_dim: 256 + num_heads: 8 + dropout: 0.1 + residual: true + +node_features_dim: 22 +edge_features_dim: 22 + +heads: + edge_classification: + enabled: true + hidden_dim: 128 + graph_classification: + enabled: true + num_classes: 4 + joint_type: + enabled: true + num_classes: 12 + dof_regression: + enabled: true + dof_tracking: + enabled: true diff --git a/configs/training/finetune.yaml b/configs/training/finetune.yaml new file mode 100644 index 0000000..09e8bae --- /dev/null +++ b/configs/training/finetune.yaml @@ -0,0 +1,45 @@ +# Fine-tuning on real data config +phase: finetune + +dataset: fusion360 +model: baseline + +pretrained_checkpoint: checkpoints/pretrain/best_val_loss.ckpt + +optimizer: + name: adamw + lr: 1e-5 + weight_decay: 1e-4 + +scheduler: + name: cosine_annealing + T_max: 50 + eta_min: 1e-7 + +training: + epochs: 50 + batch_size: 32 + gradient_clip: 1.0 + early_stopping_patience: 10 + amp: true + freeze_encoder: false # set true for frozen encoder experiment + +loss: + edge_weight: 1.0 + graph_weight: 0.5 + joint_type_weight: 0.3 + dof_weight: 0.2 + redundant_penalty: 2.0 + +checkpointing: + save_best_val_loss: true + save_best_val_accuracy: true + save_every_n_epochs: 5 + checkpoint_dir: checkpoints/finetune + +logging: + backend: wandb + project: kindred-solver + log_every_n_steps: 20 + +seed: 42 diff --git a/configs/training/pretrain.yaml b/configs/training/pretrain.yaml new file mode 100644 index 0000000..90cd232 --- /dev/null +++ b/configs/training/pretrain.yaml @@ -0,0 +1,42 @@ +# Synthetic pre-training config +phase: pretrain + +dataset: synthetic +model: baseline + +optimizer: + name: adamw + lr: 1e-3 + weight_decay: 1e-4 + +scheduler: + name: cosine_annealing + T_max: 100 + eta_min: 1e-6 + +training: + epochs: 100 + batch_size: 64 + gradient_clip: 1.0 + early_stopping_patience: 10 + amp: true + +loss: + edge_weight: 1.0 + graph_weight: 0.5 + joint_type_weight: 0.3 + dof_weight: 0.2 + redundant_penalty: 2.0 # safety loss multiplier + +checkpointing: + save_best_val_loss: true + save_best_val_accuracy: true + save_every_n_epochs: 10 + checkpoint_dir: checkpoints/pretrain + +logging: + backend: wandb # or tensorboard + project: kindred-solver + log_every_n_steps: 50 + +seed: 42 diff --git a/data/fusion360/.gitkeep b/data/fusion360/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/data/processed/.gitkeep b/data/processed/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/data/splits/.gitkeep b/data/splits/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/data/synthetic/.gitkeep b/data/synthetic/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..dfe8b4c --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,39 @@ +services: + train: + build: + context: . + dockerfile: Dockerfile + target: base + volumes: + - .:/workspace + - ./data:/workspace/data + deploy: + resources: + reservations: + devices: + - driver: nvidia + count: all + capabilities: [gpu] + command: make train + environment: + - CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0} + - WANDB_API_KEY=${WANDB_API_KEY:-} + + test: + build: + context: . + dockerfile: Dockerfile + target: cpu-only + volumes: + - .:/workspace + command: make check + + data-gen: + build: + context: . + dockerfile: Dockerfile + target: base + volumes: + - .:/workspace + - ./data:/workspace/data + command: make data-gen diff --git a/docs/.gitkeep b/docs/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/export/.gitkeep b/export/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/freecad/__init__.py b/freecad/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/freecad/bridge/__init__.py b/freecad/bridge/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/freecad/tests/__init__.py b/freecad/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/freecad/workbench/__init__.py b/freecad/workbench/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..439159d --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,97 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[project] +name = "kindred-solver" +version = "0.1.0" +description = "Assembly constraint prediction via GNN for Kindred Create" +readme = "README.md" +license = "Apache-2.0" +requires-python = ">=3.11" +authors = [ + { name = "Kindred Systems" }, +] +classifiers = [ + "Development Status :: 3 - Alpha", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Topic :: Scientific/Engineering", +] +dependencies = [ + "torch>=2.2", + "torch-geometric>=2.5", + "numpy>=1.26", + "scipy>=1.12", +] + +[project.optional-dependencies] +train = [ + "wandb>=0.16", + "tensorboard>=2.16", + "hydra-core>=1.3", + "omegaconf>=2.3", + "matplotlib>=3.8", + "networkx>=3.2", +] +freecad = [ + "pyside6>=6.6", +] +dev = [ + "pytest>=8.0", + "pytest-cov>=4.1", + "ruff>=0.3", + "mypy>=1.8", + "pre-commit>=3.6", +] + +[project.urls] +Repository = "https://git.kindred-systems.com/kindred/solver" + +[tool.hatch.build.targets.wheel] +packages = ["solver", "freecad"] + +[tool.ruff] +target-version = "py311" +line-length = 100 + +[tool.ruff.lint] +select = [ + "E", # pycodestyle errors + "W", # pycodestyle warnings + "F", # pyflakes + "I", # isort + "N", # pep8-naming + "UP", # pyupgrade + "B", # flake8-bugbear + "SIM", # flake8-simplify + "TCH", # flake8-type-checking + "RUF", # ruff-specific +] + +[tool.ruff.lint.isort] +known-first-party = ["solver", "freecad"] + +[tool.mypy] +python_version = "3.11" +warn_return_any = true +warn_unused_configs = true +disallow_untyped_defs = true +check_untyped_defs = true + +[[tool.mypy.overrides]] +module = [ + "torch.*", + "torch_geometric.*", + "scipy.*", + "wandb.*", + "hydra.*", + "omegaconf.*", +] +ignore_missing_imports = true + +[tool.pytest.ini_options] +testpaths = ["tests", "freecad/tests"] +addopts = "-v --tb=short" diff --git a/solver/__init__.py b/solver/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/solver/datagen/__init__.py b/solver/datagen/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/solver/datasets/__init__.py b/solver/datasets/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/solver/evaluation/__init__.py b/solver/evaluation/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/solver/inference/__init__.py b/solver/inference/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/solver/models/__init__.py b/solver/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/solver/training/__init__.py b/solver/training/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29